"use strict";
/* simulation parameters */
const NMESH = 512; // mesh points
const DT = 0.1; // time step
const NSTEP = 50000; // number of steps
const NECAL = 1;
const NNCAL = 1000;

/* system parameters */
const M = 2000; // mass
const LX = 50; // box length
const initCond = {
  x: 19, //position
  spread: 1, // spread
  p: 20, // momentum
};

let dx = LX / NMESH;

let hamEl = (x) => {
  let hmat = [[],[]];
  let A=0.01;
  let B=1.6;
  let C=0.005;  //SB: 0.005
  let D=1.0;

  x = x-0.5*LX;

  hmat[0][0] = x!==0 ? A*(1-Math.exp(-B*Math.abs(x)))*x/Math.abs(x) : 0;
  hmat[0][1] = C*Math.exp(-D*x*x);
  hmat[1][0] = hmat[0][1];
  hmat[1][1] = -hmat[0][0];

  return hmat;
}

let ergEl = (hmat) => {
  let erg = [];


  let discr = Math.sqrt((hmat[0][0] - hmat[1][1])*(hmat[0][0] - hmat[1][1])+4*hmat[0][1]*hmat[1][0]);
  erg[0] = 0.5*(hmat[0][0] + hmat[1][1] + discr);
  erg[1] = 0.5*(hmat[0][0] + hmat[1][1] - discr);

  return erg;
}

const EIGVEC = (x) => {
  let hmat = hamEl(x);
  return [[ hmat[0][1] / (ergEl(hmat)[0] - hmat[0][0]),
    1
  ],
  [ hmat[0][1] / (ergEl(hmat)[1] - hmat[0][0]),
    1
  ]]
}

let initWavefn = (x) => {

  if (x > LX) {
    return initWavefn(x - LX);
  }

  let gauss = Math.exp(-(x-initCond.x)*(x-initCond.x)/4.0/(initCond.spread*initCond.spread));

  let comps = [[],[]];

  comps[0][0] = gauss*Math.cos(initCond.p*(x-initCond.x));
  comps[0][1] = gauss*Math.sin(initCond.p*(x-initCond.x));
  comps[1][0] = 0;
  comps[1][1] = 0;

  return comps;
}

let computeNormFac = (wavefn) => {
  let sumsq = 0;
  for (let i = 0; i < NMESH; i++) {
    let wf = wavefn(i*dx)
    sumsq += wf[0]*wf[0]+wf[1]*wf[1];
  }

  if (sumsq > 0) {
    return 1/Math.sqrt(sumsq);
  } else if (sumsq === 0) {
    return 1;
  } else {
    throw new Error('noooooo!');
  }
}

let normFac = [computeNormFac((x)=>initWavefn(x)[0]),
computeNormFac((x)=>initWavefn(x)[1])];

let normWavefn = (x) => {
  return [initWavefn(x)[0].map((v)=>v*normFac[0]),initWavefn(x)[1].map((v)=>v*normFac[1])];
}

let vverlet = (x, p, pot) => {
  let vel = p/M;

  let force = (x) => {
    if (x >= dx) return - 0.5*(pot(x+dx)-pot(x-dx))/dx;
    else return 0;
  }

  let xnew = x + vel*DT+0.5*force(x)/M*DT*DT;

  let fnew = force(xnew);
  let vnew = vel + 0.5*(fnew+force(x))/M*DT;
  let pnew = vnew*M;

  return [xnew, pnew, vnew]
}

let matD = (s1, s2) => {
  let termD = 0;
  let diffS2 = (x) => {
    if (x > dx)
      return 0.5*(EIGVEC(x+dx)[s2] - EIGVEC(x-dx)[s2]) / dx;

    return (EIGVEC(x+dx)[s2] - EIGVEC(x)[s2]) / dx;
  }
  return (x) => EIGVEC(x)[s1]*diffS2(x);
};

matD = (s1, s2) => {
  let completeWf = [];
  completeWf[0] =

  return (x) => {

  }
}

let singleStep = ({x, p, state, coeffs}) => {
  let potcurr = (xv) => {return ergEl(hamEl(xv))[state]};
  let potother = (xv) => {return ergEl(hamEl(xv))[+!state]};
  let [xnew, pnew, vel] = vverlet(x, p, potcurr);


  let coeffsnew = coeffs;

  let matA = (s1, s2) => {
    return coeffsnew[s1]*coeffsnew[s2]
  };

  let matB = (s1, s2) => {
    return -2*matA(s2,s1)*matD(s1,s2)(xnew)*vel;
  };

  let probg = DT*matB(state, +!state) / matA(state, state); // g12

  let ifHop = Math.random() < probg;

  if (ifHop && (pnew*pnew*0.5/M < (potother(xnew) - potcurr(xnew)))) {
    ifHop = false;
  }

  if (ifHop) {
    pnew = Math.sqrt(pnew*pnew+2*M*(potcurr(xnew) - potother(xnew)));
  }
  return {xnew, pnew, ifHop, coeffsnew};
}

/*----------------*/

let coeffs = normWavefn(initCond.x)[0];

initCond.state = 0;
initCond.coeffs = normWavefn(initCond.x)[initCond.state];

let currCond = initCond;

for (let i = 0; i < 5; i++) {
  let newCond = {};
  let res = singleStep(currCond);
  //console.log(res);

  [newCond.x, newCond.p] = [res.xnew, res.pnew];

  let newNormWavefn;

  if (res.ifHop) {
    newCond.state = +!currCond.state;
    newCond.coeffs[newCond.state] = res.coeffsnew[currCond.state];
    newCond.coeffs[currCond.state] = [0,0];
  } else {
    newCond.state = currCond.state;
    newCond.coeffs = res.coeffsnew;
  }

  currCond = newCond;
}

for (let i = 0; i < 1; i++) {
  console.log((matD(0,1))(2));
}