#include<iostream>
#include<fstream>
#include<cmath>
#include<ctime>
#include<vector>
#include<cstdlib>
using namespace std;

// global variables
long double alpha = sqrt(8.0);
int gam = 1;
long double deltax = .01;
long double deltaeps = .00001;
long double Tg = 1;
long double r0 = 1;  
ofstream fout1, fout2, fout3, fout4, fout5, fout6, fout7, fout8, fout9, fout10;

// functions -- note "x" in this code == "u" in the published paper.
long double rhoval(long double x);
long double kxfac(long double x);
void openfiles(long double T, long double doteps);

int main(int argc, const char *argv[]) {
// run timers
   time_t tstart, tend, tmid;
   tstart = time(0);
// key variables
   long double T = atof(argv[1]);
   long double doteps = atof(argv[2]);
   long double maxeps = atof(argv[3]) + deltaeps;
   long double deltat = deltaeps/doteps;
   int nisteps = alpha*alpha/deltax;
   int i, j, k, u, v;    // i -- x-zone depth, j = strain, k = time
   int DI, DJ;  // binning for ELP output
   DI = 10;  DJ = 100;
   long double fac2, x, thiseps, eps, sigma, sigmael, sigmapl, sigmaentr, Ztot, Ztilde;
   long double dE, dS, E, S, F, W, E0, S0, F0, Slast;  // thermodynamic quantities
   long double xavg, epsavg, zthis, sigmathis, tauthis;
   long double totoutflow, totinflow, share, peq, fboltz, sharedenom, sumz;
   long double tauinv, tauav, tausqavg;
   long double ssfac = alpha*alpha*alpha*alpha;

    cout << "Rob Hoy's thermal plasticity code" << endl;
    cout << "If you use results from this code in a publication, please cite: R. S. Hoy, Physical Review E v. 96, 063001 (2017)." << endl;
    
// set ell[i], ellmax, njsteps
   vector <int> ell;
   ell.resize(nisteps);
   int ellmax;
   for (i = 0; i < nisteps; i++) {
      x = deltax*(i+0.5);
      ell[i] = static_cast<int> (ceil(1.0/(sqrt(kxfac(x))*deltaeps)));
   }
   ellmax = ell[nisteps - 1];
   int njsteps = (maxeps/deltaeps) + 2*ellmax;
 
// declare main arrays
   vector<long double> rho, Kofx, fofx;
   vector < vector<long double> > zofx, zlast, zeq;   // "z" in this code is "w" in the published paper
   vector < vector<long double> > tauofx, sigmaelofx;
   long double poftau[1200], pofsigma[1200];
   int sigmabin, taubin;

// allocate 1D arrays
   ell.resize(nisteps);		rho.resize(nisteps);           Kofx.resize(nisteps);        fofx.resize(nisteps);

// allocate 2D arrays
   cout << "Allocating arrays: isteps = " << nisteps << ", jsteps = " << njsteps << endl;
   zofx.resize(nisteps);       	zlast.resize(nisteps); 		zeq.resize(nisteps);    tauofx.resize(nisteps);
   for (i = 0; i < nisteps; i++) {
      zofx[i].resize(njsteps);          zlast[i].resize(njsteps);	    zeq[i].resize(njsteps);     tauofx[i].resize(njsteps);
   }
   cout << "Arrays allocated." << endl;

// open log files
   openfiles(T,doteps);

// set rho(x), K(x), ell(x)
   for (i = 0; i < nisteps; i++) {
      x = deltax*(i+0.5);
      Kofx[i] = 2*Tg*x*kxfac(x);
      rho[i] = rhoval(x);
   }

// set pofx, zofx
   Ztilde = 0;
   for (i = 0; i < nisteps; i++) {
      x = deltax*(i+0.5);
      fac2 = (Tg/T)*(x-alpha*alpha);
      for (j = 0; j <  2*ell[i]+1; j++) {
         thiseps = (j-ell[i])*deltaeps;
         fboltz = expl(fac2)*expl(-Kofx[i]*thiseps*thiseps/(2*T));  // Boltzmann factor
         zofx[i][ellmax-ell[i]+j] = rho[i]*fboltz;
         Ztilde += zofx[i][ellmax-ell[i]+j];
      }
   }
   cout << "Z = " << Ztilde << endl;  // this is the PARTITION function
// rescale z, set zinit & zeq
   Ztot = 0;
   for (i = 0; i < nisteps; i++) for (j = 0; j <  2*ell[i]+1; j++) {
      zofx[i][ellmax-ell[i]+j] /= Ztilde;
      zeq[i][ellmax-ell[i]+j] = zofx[i][ellmax-ell[i]+j];
      Ztot += zofx[i][ellmax-ell[i]+j];
   }
   cout << "Ztot = " << Ztot << endl; 
// now z = zinit = rho*peq/Ztilde is my "w"

// define array of characteristic stresses, relaxation times
   cout << "Setting characteristic stresses, relaxation times" << endl;
   for (i = 0; i < nisteps; i++) {
      x = deltax*(i+0.5);
      for (j = 0; j < njsteps-ellmax+ell[i]; j++) {
         thiseps = (j-ell[i])*deltaeps;
         tauofx[i][ellmax-ell[i]+j] = (1.0/r0)*expl((Tg/T)*x)*expl(-Kofx[i]*thiseps*thiseps/(2*T));  
      }
   }

// main loop: Eq. 12 of continuous-SGR paper
   W = 0;
   for (k = 0; k < maxeps/deltaeps; k++) {
// update thermodynamics:  output all quantities scaled by (k_B T_g)
      E = 0; S = 0; F = 0;
      for (i = 0; i < nisteps; i++) {
         x = deltax*(i+0.5);
         for (j = 0; j < 2*ell[i]+k+1; j++) {
// break if zone is so unlikely to be occupied that numerical overflow sets p(x, eps^el) = 0
            if (zofx[i][ellmax-ell[i]+j]/rho[i] == 0) continue;
// statistical definition of entropy: note has (deltax*deltaeps) - dependence
         dS = -logl(zofx[i][ellmax-ell[i]+j]/rho[i])*zofx[i][ellmax-ell[i]+j];
// thermodynamic-consistency check
            if (isnan(dS) || dS < 0) {
               cout << "Fuckup: k = " << k << " i = " << i << " ell[i] = " << ell[i] << " j = " << j << " Ztot = " << Ztot << " rho[i] = " << rho[i] << " z = " << zofx[i][ellmax-ell[i]+j] << " w = " << zofx[i][ellmax-ell[i]+j]/Ztot << " p = " << (zofx[i][ellmax-ell[i]+j]/rho[i])/Ztot << " dS = " << dS << endl;
              exit(0);
            }
            S += (T/Tg)*dS;
            thiseps = (j-ell[i])*deltaeps;
            dE = (kxfac(x)*thiseps*thiseps - 1.0)*x*zofx[i][ellmax-ell[i]+j]/Ztot;
            E += dE;
         }
      }
      F = E-S;
      W += sigma*deltaeps;
      if (k == 0) { E0 = E;   S0 = S;   F0 = F; }

// set current z to zlast so can properly update z
      for (i = 0; i < nisteps; i++) for (j = 0; j < 2*ell[i]+k+1; j++) zlast[i][ellmax-ell[i]+j] = zofx[i][ellmax-ell[i]+j];
// set <tau^-1> (Eq. 10)
      tauinv = 0;   tauav = 0;    tausqavg = 0;
      for (i = 0; i < nisteps; i++) for (j = 0; j < 2*ell[i]+k+1; j++) {
         tauinv += (zlast[i][ellmax-ell[i]+j]/Ztot)/tauofx[i][ellmax-ell[i]+j];
// output in terms of tau/tau0
         tauav += zlast[i][ellmax-ell[i]+j]*log10(r0*tauofx[i][ellmax-ell[i]+j]);
         tausqavg += zlast[i][ellmax-ell[i]+j]*log10(r0*tauofx[i][ellmax-ell[i]+j])*log10(r0*tauofx[i][ellmax-ell[i]+j]);
// alternative non-logarithmic averaging
//         tauav += zlast[i][ellmax-ell[i]+j]*tauofx[i][ellmax-ell[i]+j];
//         tausqavg += zlast[i][ellmax-ell[i]+j]*tauofx[i][ellmax-ell[i]+j]*tauofx[i][ellmax-ell[i]+j];
      }

// update z values
// convective term including yielding
      totoutflow = 0;
      if (k > 0) for (i = 0; i < nisteps; i++) for (j = 1; j < 2*ell[i]+k+2; j++) {
         zofx[i][ellmax-ell[i]+j] = zlast[i][ellmax-ell[i]+j-1]*expl(-deltat/tauofx[i][ellmax-ell[i]+j-1]);
         totoutflow += zlast[i][ellmax-ell[i]+j-1]*(1.0-expl(-deltat/tauofx[i][ellmax-ell[i]+j-1]));
       }
// inflow term (rhs of Eq. 12)
      if (k > 0) for (i = 0; i < nisteps; i++) for (j = 0; j < 2*ell[i]+1; j++) {
//  note this equation has no stationary solution for w(u, eps^el) since it's based on the trap model
            zofx[i][ellmax-ell[i]+j] += zeq[i][ellmax-ell[i]+j]*totoutflow; 
         }
      
// rescale z to correct for fact that sumz undergoes small changes due to finite-timestep errors
   sumz = 0;
   for (i = 0; i < nisteps; i++) for (j = 0; j < 2*ell[i]+k+1; j++) sumz += zofx[i][ellmax-ell[i]+j];
   for (i = 0; i < nisteps; i++) for (j = 0; j < 2*ell[i]+k+1; j++) zofx[i][ellmax-ell[i]+j] *= Ztot/sumz;
// track progress of the run
   tmid = time(0);
   if (!(k%10)) cout << "Step " << k << " , Ztot = " << sumz << ", tauinv = " << tauinv << ", totoutflow = " << totoutflow << ", totinflow = " << totinflow << ", time = " <<  difftime(tmid,tstart) <<  endl;

// calculate stresses
      sigma = 0;        sigmael = 0;            sigmapl = 0;
// elastic and plastic components
      for (i = 0; i < nisteps; i++) for (j = 0; j < 2*ell[i]+k+1; j++) {
         thiseps = (j-ell[i])*deltaeps;
         sigmael += zofx[i][ellmax-ell[i]+j]*Kofx[i]*thiseps;
// below 'plastic' stress was defined in the 1998 P. Sollich SGR paper but isn't sensible since is nonzero in undeformed systems
         sigmapl += zofx[i][ellmax-ell[i]+j]*Kofx[i]*(thiseps*thiseps/2)*(1-expl(-deltat/tauofx[i][ellmax-ell[i]+j]))/doteps;
      }
// entropic component
      if (k > 0) sigmaentr = -(S-Slast)/deltaeps;
      sigma = sigmael + sigmaentr;

// output current state
      if (!(k%10)) {
         cout << k*deltaeps << " " << sigma << " " << sigmael << " " << sigmapl << " " << sigmaentr << " " << tauav << " " << sqrt(tausqavg - tauav*tauav) << endl;
         fout1 << k*deltaeps << " " << sigma << endl;
         fout3 << k*deltaeps << " " << tauav << endl;
         fout4 << k*deltaeps << " " << E << " " << S << " " << F << " " << W << " " << E - E0 << " " << S - S0 << " " << F - F0 << endl;
         fout5 << k*deltaeps << " " << sqrt(tausqavg - tauav*tauav) << endl;
         fout6 << k*deltaeps << " " << sigmael << endl;
         fout8 << k*deltaeps << " " << sigmaentr << endl;
      }

// output EL position
      xavg = 0;		epsavg = 0;		zthis = 0;
      if (!(k%10)) for (i = 0; i < nisteps/DI; i++) for (j = 0; j < njsteps/DJ; j++) {
         for (u = 0; u < DI; u++) for (v = 0; v < DJ; v++) {
             x = (DI*i+u+0.5)*deltax;
             thiseps = (DJ*j+v-ellmax)*deltaeps;
             xavg += x/(DJ*DI);
             epsavg += thiseps/(DJ*DI);
             zthis += zofx[DI*i+u][DJ*j+v];
         }
         fout2 << xavg << " " << epsavg << " " << zthis/Ztot << endl;
         xavg = 0;	epsavg = 0;   zthis = 0;         
      }

// output P(sigma), P(tau)
   if (!(k%10)) {
      for (i = 0; i < 1200; i++) {  pofsigma[i] = 0;   poftau[i] = 0; }
      for (i = 0; i < nisteps; i++) for (j = 0; j < 2*ell[i]+k+1; j++) {
         thiseps = (j-ell[i])*deltaeps;
         sigmathis = Kofx[i]*thiseps/ssfac;
         tauthis = (1/2.302585093)*((1/T)*x - Kofx[i]*thiseps*thiseps/(2*T));         
         sigmabin = floor((sigmathis+3)/.01);
         taubin = floor((tauthis+5)/.05);
         if (sigmabin > -1 && sigmabin < 1200) pofsigma[sigmabin] += zofx[i][ellmax-ell[i]+j]/Ztot;
         if (taubin > -1 && taubin < 1200) poftau[taubin] += zofx[i][ellmax-ell[i]+j]/Ztot;
      }
      for (i = 0; i < 1200; i++) {
         fout9 << i*.01 - 3 << " " << pofsigma[i] << endl;
         fout10 << i*.05 - 5 << " " << poftau[i] << endl;
      }
   }
      Slast = S;
  } // end of main loop

// run time
   tend = time(0);
   cout << "Execution took " << difftime(tend, tstart) << " seconds." <<  endl;

   fout1.close();       fout2.close();          fout3.close();
   fout4.close();       fout5.close();          fout6.close();
   fout8.close();	fout9.close();		fout10.close();

   return 0; 
}

long double kxfac(long double x) {
   return 400;  // this is k(u) in the paper
}


// density of states rho(x)
long double rhoval(long double x) {
   long double prefac, cutfac, expfac, divfac, w;
   prefac = (1.0/alpha)*powl(x/alpha,gam-1);
   cutfac = 1.0 - powl(x/(alpha*alpha),gam);
   expfac = expl(-powl(x/alpha,gam)/gam);
   divfac = 1 + gam*powl(alpha,-gam)*(expl(-pow(alpha,gam)/gam)-1);
   w = prefac*cutfac*expfac/divfac;
   return w;
}

// set up output files: sigma, EL-position, tau, <x>
void openfiles(long double T, long double doteps) {
   char ofname[50];
   sprintf(ofname,"sigmaofepsilon.gamma%d.T%.3f.doteps%.6f",gam,static_cast<double>(T),static_cast<double>(doteps));
   fout1.open(ofname);
   sprintf(ofname,"ELpositionofx.gamma%d.T%.3f.doteps%.6f",gam,static_cast<double>(T),static_cast<double>(doteps));
   fout2.open(ofname);
   sprintf(ofname,"tauofepsilon.gamma%d.T%.3f.doteps%.6f",gam,static_cast<double>(T),static_cast<double>(doteps));
   fout3.open(ofname);
   sprintf(ofname,"ESFofepsilon.gamma%d.T%.3f.doteps%.6f",gam,static_cast<double>(T),static_cast<double>(doteps));
   fout4.open(ofname);
   sprintf(ofname,"deltatauofepsilon.gamma%d.T%.3f.doteps%.6f",gam,static_cast<double>(T),static_cast<double>(doteps));
   fout5.open(ofname);
   sprintf(ofname,"sigmaelofepsilon.gamma%d.T%.3f.doteps%.6f",gam,static_cast<double>(T),static_cast<double>(doteps));
   fout6.open(ofname);
   sprintf(ofname,"sigmaentrofepsilon.gamma%d.T%.3f.doteps%.6f",gam,static_cast<double>(T),static_cast<double>(doteps));
   fout8.open(ofname); 
   sprintf(ofname,"Pofsigmaelofepsilon.gamma%d.T%.3f.doteps%.6f",gam,static_cast<double>(T),static_cast<double>(doteps));
   fout9.open(ofname);
   sprintf(ofname,"Poftauofepsilon.gamma%d.T%.3f.doteps%.6f",gam,static_cast<double>(T),static_cast<double>(doteps));
   fout10.open(ofname);  
   return;
}