// generic 2D physics environment, with support for lateral and vertical advection and mixing,
// sinking, variable light profiles, and boundary inputs.
// this + a method for calculating the velocity field = tonicella.

class Slice extends Environment {
  
  int M,K;
  float[] dx, dy, dz, z_rho, z_w, x_rho, x_u;
  int[] kbot, kbotu;
  float[][] Axy, Ayz, V;
  float[][] u, w, kappa;
  float[] KH;
  float[] sinkingRate, benthicLossFrac, benthicRecyclingFrac;
  BoundaryCondition leftBC = new BoundaryCondition(this);
  BoundaryCondition rightBC = new BoundaryCondition(this);
  BoundaryCondition[] lateralInputs = new BoundaryCondition[0];
  int[] lateralInputM = new int[0];
  
  Slice() {}
  
  Slice(Ecosystem eco, int M, int K, float Ltot, float Htot) {
    createDomain(M,K,Ltot,Htot);
    linkToEcosystem(eco);
    addParam("surface PAR","PAR0","W/m2", 150, 0, 300);
    allocateStorageForPhysics();
  }
    
  Slice(Ecosystem eco, int M, int K, float Ltot, float Hhead, float Hmouth, float Bhead, float Bmouth) {
    createDomain(M,K,Ltot,Hhead,Hmouth,Bhead,Bmouth);
    linkToEcosystem(eco);
    addParam("surface PAR","PAR0","W/m2", 150, 0, 300);
    allocateStorageForPhysics();
  }

  void allocateStorageForPhysics() {
    // define physical fluxes
    for (int vi=0; vi<vars.length; vi++) {
      if (vars[vi].length==Ncells) {
        if (vars[vi].advects) {
          addFlux("horizontal advection of " + vars[vi].shortname, "hAdv-" + vars[vi].shortname, null, vars[vi], Ncells);
          addFlux("vertical advection of " + vars[vi].shortname, "vAdv-" + vars[vi].shortname, null, vars[vi], Ncells); 
          addFlux("horizontal diffusion of " + vars[vi].shortname, "hDiff-" + vars[vi].shortname, null, vars[vi], Ncells); // horizontal diffusion is considered a special kind of advection
        }
        if (vars[vi].mixes) {
          addFlux("vertical diffusion of" + vars[vi].shortname, "diff-" + vars[vi].shortname, null, vars[vi], Ncells);
        }
        if (vars[vi].sinks) { // if sinking is turned on later on (as it probably will be), this flux will be created then
          addFlux("sinking of" + vars[vi].shortname, "sink-" + vars[vi].shortname, null, vars[vi], Ncells);      
        }
        addFlux("lateral input of " + vars[vi].shortname, "lateral-" + vars[vi].shortname, null, vars[vi], Ncells);
      }
    }
    sinkingRate = new float[vars.length];  
    benthicLossFrac = new float[vars.length];
    benthicRecyclingFrac = new float[vars.length];
    addVar("horizontal velocity","u","m/day",Ncells);
    addVar("vertical velocity","w","m/day",Ncells);
  }
  
  
  void startArchive(float dt) {
    super.startArchive(dt);
    if (archive==null) return;
    archive.addNameValuePair("M",M);
    archive.addNameValuePair("K",K);
    archive.addNameValuePair("dx",dx);
    archive.addNameValuePair("dy",dy);
    archive.addNameValuePair("dz",dz);
    archive.addNameValuePair("Axy",Axy);
    archive.addNameValuePair("Ayz",Ayz);
    archive.addNameValuePair("V",V);
    archive.addNameValuePair("x_rho",x_rho);
    archive.addNameValuePair("x_u",x_u);
    archive.addNameValuePair("z_rho",z_rho);
    archive.addNameValuePair("z_w",z_w);
    float[] kbotf = new float[M];
    for (int i=0; i<kbot.length; i++) kbotf[i] = kbot[i]; // int -> float
    archive.addNameValuePair("kbot",kbotf);
    float[] kbotuf = new float[M+1];
    for (int i=0; i<kbotu.length; i++) kbotuf[i] = kbotu[i]; // int -> float
    archive.addNameValuePair("kbotu",kbotuf);
  }
  
  
  
  // geometry ------------------------------------------------------------------------
  
  
  
  void createDomain(int M, int K, float Ltot, float Htot) {
    dx = new float[M];
    for (int m=0; m<M; m++) dx[m] = Ltot/M;
    dz = new float[K];
    for (int k=0; k<K; k++) dz[k] = Htot/K;
    createDomain(dx,dz);
  }
  
  void createDomain(int M, int K, float Ltot, float Hhead, float Hmouth, float Bhead, float Bmouth) {
    createDomain(M,K,Ltot,max(Hhead, Hmouth));
    float[] B = new float[M+1];
    for (int m=0; m<M+1; m++) B[m] = map(m, 0, M, Bmouth, Bhead);
    setWidths(B);
    float[] H = new float[M];
    for (int m=0; m<M; m++) H[m] = map(m, -0.5, M-0.5, -Hmouth, -Hhead);
    setDepths(H);      
  }
    
  void createDomain(float[] dx, float[] dz) {
    M = dx.length;
    K = dz.length;
    Ncells = M*K;
    
    // lengths
    this.dx = dx;
    x_rho = new float[M];
    x_u = new float[M+1];
    x_u[0] = 0;
    for (int m=0; m<M; m++) {
      x_u[m+1] = x_u[m] + dx[m];
      x_rho[m] = x_u[m] + dx[m]/2;
    }
    
    // depths
    this.dz = dz;
    z_rho = new float[K];
    z_w = new float[K+1];
    z_w[K] = 0;
    for (int k=K-1; k>=0; k--) {
      z_w[k] = z_w[k+1] - dz[k];
      z_rho[k] = z_w[k+1] - dz[k]/2;
    }
    kbot = new int[M];
    for (int m=0; m<M; m++) kbot[m] = 0;
    kbotu = new int[M+1];
    for (int m=0; m<M+1; m++) kbotu[m] = 0;
    
    // widths
    setWidths(1);

    // initialize internal parameters
    u = new float[K][M+1];
    w = new float[K+1][M];
    kappa = new float[K+1][M];
    KH = new float[M+1];
    for (int m=0; m<M; m++) {
      for (int k=0; k<K+1; k++) w[k][m] = 0;
      for (int k=0; k<K+1; k++) kappa[k][m] = 0;
    }
    for (int m=0; m<M+1; m++) {
      for (int k=0; k<K; k++) u[k][m] = 0;
      KH[m] = 0;
    }
  }
  
/*
  void createDomain(String filename) { // load everything from a netcdf file. This could be extended to deal gracefully with missing variables
    NetcdfFile nc = nc_open(filename);
    x_rho = nc_read1D(nc,"x_rho");
    x_u = nc_read1D(nc,"x_u");
    M = x_rho.length;
    dx = new float[M];
    for (int m=0; m<M; m++) dx[m] = x_u[m+1] - x_u[m];
    
    z_rho = nc_read1D(nc,"z_rho");
    z_w = nc_read1D(nc,"z_w");
    K = z_rho.length;
    kbot = round(nc_read1D(nc,"kbot"));
    dz = new float[K];
    for (int k=0; k<K; k++) dz[k] = z_w[k+1] - z_w[k];
    kbotu = new int[M+1];
    kbotu[0] = kbot[0];
    for (int m=1; m<M; m++) kbotu[m] = max(kbot[m-1],kbot[m]);
    kbotu[M] = kbot[M-1];
    
    dy = nc_read1D(nc,"dy");
    setWidths(dy);

    nc_close(nc);
    
    Ncells = M*K;
    // initialize internal parameters
    u = new float[K][M+1];
    w = new float[K+1][M];
    kappa = new float[K+1][M];
    KH = new float[M+1];
    for (int m=0; m<M; m++) {
      for (int k=0; k<K+1; k++) w[k][m] = 0;
      for (int k=0; k<K+1; k++) kappa[k][m] = 0;
    }
    for (int m=0; m<M+1; m++) {
      for (int k=0; k<K; k++) u[k][m] = 0;
      KH[m] = 0;
    }
  }
*/
  
  void setDepths(float[] zbot) { // once the domain has been defined, set channel depths in each reach via the index kbot
    if (zbot.length != M) println("setDepths(): zbot should be length " + M + ".");
    for (int m=0; m<M; m++) {
      while (kbot[m] < K+1 && z_w[kbot[m]] < zbot[m]) kbot[m]++;
    }
    kbotu = new int[M+1];
    kbotu[0] = kbot[0];
    for (int m=1; m<M; m++) kbotu[m] = max(kbot[m-1],kbot[m]);
    kbotu[M] = kbot[M-1];
  }
  
  void setWidths(float dy0) {
    dy = new float[M+1];
    for (int m=0; m<M+1; m++) dy[m] = dy0;
    setWidths(dy);
  }
  
  void setWidths(float[] dy) { // set widths directly, calculate areas and volumes accordingly
    this.dy = dy;
    Axy = new float[K+1][M];
    Ayz = new float[K][M+1];
    V = new float[K][M];
    for (int m=0; m<M; m++) {
      for (int k=0; k<K+1; k++) {
        Axy[k][m] = 0.5 * (dy[m] + dy[m+1]) * dx[m];
      }
    }
    for (int m=0; m<M+1; m++) {
      for (int k=0; k<K; k++) {
        Ayz[k][m] = dy[m] * dz[k];
      }
    }
    for (int m=0; m<M; m++) {        
      for (int k=0; k<K; k++) {
        V[k][m] = Axy[k][m] * dz[k];
      }
    }
  }
  
  
  // utility ------------------------------------------------------------------------
  
  
  
  int km2ind(int k, int m) {
    return m*K + k;
  }
  
  int ind2m(int ind) {
    return ind / K;
  }
  
  int ind2k(int ind) {
    return ind % K;
  }
  
  int[] ind2km(int ind) {
    return new int[] {ind2k(ind), ind2m(ind)};
  }

  
  // top-level flux calculations ------------------------------------------------------
  
  
  
  void calcFluxes() {
    calcVelocity();
    doHAdvection();
    doVAdvection();
    doHDiffusion();
    doMixing();
    doLateralInputs();
    doSinking();
    calcLight();
    
    eco.calcFluxes();
    
    // store a copy of u and w in statevars, so that they get saved to the archive.
    // u[:][M] and w[K][:] get omitted, so that these statevars are a normal size.
    // a less klugey way to do this would be to define Statevars of arbitrary size that can handle their own k,m <-> ind conversion internally.
    Statevar uvar = getVar("u");
    Statevar wvar = getVar("w");
    for (int m=0; m<M; m++) {
      for (int k=0; k<K; k++) {
        uvar.current[km2ind(k,m)] = u[k][m];
        wvar.current[km2ind(k,m)] = w[k][m];
      }
    }
  }
  
  
    
  // velocity calculation -------------------------------------------------------------
  
  
  void calcVelocity() {
    // if you want u to be nonzero, set it here.
    for (int m=0; m<M+1; m++) for (int k=0; k<kbotu[m]; k++) u[k][m] = 0; // make sure u below the bottom is 0
    wFromMassCons();
  }
  
  void wFromMassCons() {
    // set w via mass conservation
    for (int m=0; m<M; m++) {
      for (int k=0; k<kbot[m]+1; k++) w[k][m] = 0;
      for (int k=kbot[m]+1; k<K; k++) {
        float FH = totalHorizVolFlux(k-1,m);
        w[k][m] = (FH + Axy[k-1][m] * w[k-1][m]) / Axy[k][m];
      }
      w[K][m] = 0; // w at the surface is guaranteed 0. If the integrated horizontal fluxes don't balance, the excess will accumulate in the top cell
    }
  }
  
  float totalHorizVolFlux(int k, int m) {
    float F = u[k][m] * Ayz[k][m] - u[k][m+1] * Ayz[k][m+1];
    // add contribution from lateral inputs into cell m. (This should hold whether the input volume flux exits on the left or the right.)
    for (int ii=0; ii<lateralInputM.length; ii++) {
      if (lateralInputM[ii] == m) {
        float[] v = lateralInputs[ii].getVelocity();
        F += v[k] * dz[k] * dx[m];
      }
    }
    return F;
  }
  
  
  // horizontal advection --------------------------------------------------------------- 


  void doHAdvection() {
    for (int vi=0; vi<vars.length; vi++) {
      if (vars[vi].advects) {
        // for each advecting statevar...
        
        // get b.c.'s. If null, use adjacent grid cell: this amounts to a zero-gradient b.c.
        float[] leftConc = leftBC.getVar(vars[vi].shortname);
        if (leftConc==null) {
          leftConc = new float[K];
          for (int k=0; k<K; k++) leftConc[k] = vars[vi].current[km2ind(k,0)];
        }
        float[] rightConc = rightBC.getVar(vars[vi].shortname);
        if (rightConc==null) {
          rightConc = new float[K];
          for (int k=0; k<K; k++) rightConc[k] = vars[vi].current[km2ind(k,M-1)];
        }
        
        // extensive fluxes A * u * conc
        float[][] F = new float[K][M+1];
        for (int m=0; m<M+1; m++) {
          for (int k=kbotu[m]; k<K; k++) {        
            float upstreamConc;
            if (u[k][m] > 0) {
              if (m==0) {upstreamConc = leftConc[k];} else {upstreamConc = vars[vi].current[km2ind(k,m-1)];}
            } else {
              if (m==M) {upstreamConc = rightConc[k];} else {upstreamConc = vars[vi].current[km2ind(k,m)];}
            }
            F[k][m] = Ayz[k][m] * u[k][m] * upstreamConc; // u must be 0 below kbot, otherwise this will be nonzero
          }
        }
        
        // extensive -> intensive
        Flux theFlux = getFlux("hAdv-" + vars[vi].shortname);
        for (int m=0; m<M; m++) {
          for (int k=kbot[m]; k<K; k++) {
            theFlux.current[km2ind(k,m)] = (F[k][m] - F[k][m+1]) / V[k][m];
          }
        }
      }
    }
  }
       

  // vertical advection --------------------------------------------------------------- 
  
  
  void doVAdvection() {
    // for each column...
    for (int m=0; m<M; m++) {
      // create temporary storage for extensive fluxes
      float[] F = new float[K+1];
      for (int k=0; k<K+1; k++) F[k] = 0;
      // for each advecting statevar...
      for (int vi=0; vi<vars.length; vi++) {
        if (vars[vi].advects) {
          // calc extensive fluxes across each interface
          for (int k=kbot[m]+1; k<K; k++) {
            if (w[k][m] > 0) {
              F[k] = Axy[k][m] * w[k][m] * vars[vi].current[km2ind(k-1,m)];
            } else {
              F[k] = Axy[k][m] * w[k][m] * vars[vi].current[km2ind(k,m)];
            }
          }
          // convert to intensive fluxes for each grid cell
          Flux theFlux = getFlux("vAdv-" + vars[vi].shortname);
          for (int k=kbot[m]; k<K; k++) {
            theFlux.current[km2ind(k,m)] = (F[k] - F[k+1]) / V[k][m];
          }
        }
      }
    }
  } 
 
 
  // horizontal diffusion ------------------------------------------------------------- 


  void doHDiffusion() {
      for (int vi=0; vi<vars.length; vi++) {
        if (vars[vi].advects) {

        // get b.c.'s. If null, use adjacent grid cell: this amounts to a zero-gradient b.c.
        float[] leftConc = leftBC.getVar(vars[vi].shortname);
        if (leftConc==null) {
          leftConc = new float[K];
          for (int k=0; k<K; k++) leftConc[k] = vars[vi].current[km2ind(k,0)];
        }
        float[] rightConc = rightBC.getVar(vars[vi].shortname);
        if (rightConc==null) {
          rightConc = new float[K];
          for (int k=0; k<K; k++) rightConc[k] = vars[vi].current[km2ind(k,M-1)];
        }

        // extensive fluxes
        float[][] F = new float[K][M+1];
        for (int k=0; k<K; k++) {
          F[k][0] = - Ayz[k][0] * KH[0] * (vars[vi].current[km2ind(k,0)] - leftConc[k]) / dx[0];
          for (int m=1; m<M; m++) {
            if (k >= kbotu[m]) {
              F[k][m] = -2 * Ayz[k][m] * KH[m] * (vars[vi].current[km2ind(k,m)] - vars[vi].current[km2ind(k,m-1)]) / (dx[m] + dx[m-1]);
            }
          }
          F[k][M] = - Ayz[k][M] * KH[M] * (rightConc[k] - vars[vi].current[km2ind(k,M-1)]) / dx[M-1];
        }
          
        // extensive -> intensive
        Flux theFlux = getFlux("hDiff-" + vars[vi].shortname);
        for (int m=0; m<M; m++) {
          for (int k=kbot[m]; k<K; k++) {
            theFlux.current[km2ind(k,m)] = (F[k][m] - F[k][m+1]) / V[k][m];
          }
        }
      }
    }
  }


  // mixing --------------------------------------------------------------------------- 
  
  
  void doMixing() {
    for (int m=0; m<M; m++) {
      for (int vi=0; vi<vars.length; vi++) {
        if (vars[vi].mixes) {
          // extensive fluxes across each interface
          float[] F = new float[K+1];
          for (int k=0; k<kbot[m]+1; k++) {
            F[k] = 0;
          }
          for (int k=kbot[m]+1; k<K; k++) {
            F[k] = -2 * Axy[k][m] * kappa[k][m] * (vars[vi].current[km2ind(k,m)] - vars[vi].current[km2ind(k-1,m)]) / (dz[k] + dz[k-1]);
          }
          F[K] = 0;
          // convert to intensive fluxes for each grid cell
          Flux theFlux = getFlux("diff-" + vars[vi].shortname);
          for (int k=kbot[m]; k<K; k++) {
            theFlux.current[km2ind(k,m)] = (F[k] - F[k+1]) / V[k][m];
          }
        }
      }
    }
  }
  
  
  // lateral inputs --------------------------------------------------------------------
  
  
  BoundaryCondition addLateralInput(float x0, float Q, float h) {
    // x0 is the distance along the channel where the BC enters.
    // Q is the total volme flux.
    // the total tracer load (BC.getVar("some tracer") * Q) assumed to enter uniformly from -h .. 0.
    BoundaryCondition BC = new BoundaryCondition(this);
    lateralInputs = (BoundaryCondition[])append(lateralInputs, BC);
    int[] m01 = findBeforeAfter(x_u, x0);
    int m0 = m01[0];
    if (m0==-1) println("can't add a lateral input at " + x0 + " (valid range " + x_u[0] + " to " + x_u[M] + ")");
    lateralInputM = (int[])append(lateralInputM,m0);
    BC.setVolFlux(Q);
    BC.setVelocityFromVolFlux(m0,h);
    return BC;
  }
  
  
  void doLateralInputs() {
    // initialize
    for (int vi=0; vi<vars.length; vi++) {
      Flux theflux = getFlux("lateral-" + vars[vi].shortname);
      if (theflux != null) for (int ind=0; ind<theflux.current.length; ind++) theflux.current[ind] = 0;
    }
    // step through inputs
    for (int ii=0; ii<lateralInputs.length; ii++) {
      BoundaryCondition input = lateralInputs[ii];
      int m = lateralInputM[ii];
      float[] v = input.getVelocity();
      // for each statevar found in the lateral input, add the flux
      for (int vi=0; vi<vars.length; vi++) {
        float[] c = input.getVar(vars[vi].shortname);
        if (c != null) {
          Flux theflux = getFlux("lateral-" + vars[vi].shortname);
          for (int k=kbot[m]; k<K; k++) {
            theflux.current[km2ind(k,m)] += c[k] * v[k] * dx[m] * dz[k] / V[k][m];
          }
        }
      }
    }
  }
  
  
  // sinking --------------------------------------------------------------------------- 

  
  void addSinking(String varname, float wsink)                 {addSinking(varname, wsink, 0, 0, null);}    
  void addSinking(String varname, float wsink, float lossFrac) {addSinking(varname, wsink, lossFrac, 0, null);}
  
  void addSinking(String varname, float wsink, float lossFrac, float recyclingFrac, Statevar recyclingDest) {
    Statevar var = getVar(varname);
    var.sinks = true;
    addFlux("sinking of" + var.shortname, "sink-" + var.shortname, null, var, Ncells);      
    int i = getVarIndex(varname);
    sinkingRate[i] = wsink;
    benthicLossFrac[i] = lossFrac;
    addFlux("sedloss-" + var.shortname, "sedloss-" + var.shortname, null, var, Ncells);
    benthicRecyclingFrac[i] = recyclingFrac;
    addFlux("sedrecyc-" + var.shortname, "sedrecyc-" + var.shortname, var, recyclingDest, Ncells);
  }

  void doSinking() {
    for (int m=0; m<M; m++) {
      for (int vi=0; vi<vars.length; vi++) {
        if (vars[vi].sinks) {
          // sinking flux within the water column
          float[] F = new float[K+1];
          for (int k=0; k<kbot[m]+1; k++) {
            F[k] = 0;
          }
          for (int k=kbot[m]+1; k<K; k++) {
            F[k] = - Axy[k][m] * sinkingRate[vi] * vars[vi].current[km2ind(k,m)];
          }
          F[K] = 0;
          // extensive -> intensive
          Flux theFlux = getFlux("sink-" + vars[vi].shortname);
          for (int k=kbot[m]; k<K; k++) {
            theFlux.current[km2ind(k,m)] = (F[k] - F[k+1]) / V[k][m];
          }
          // cross-boundary sinking flux
          for (int k=0; k<kbot[m]; k++) {
            F[k] = 0;
          }
          F[kbot[m]] = - Axy[kbot[m]][m] * sinkingRate[vi] * vars[vi].current[km2ind(kbot[m],m)];
          for (int k=kbot[m]+1; k<K; k++) {
            F[k] = - (Axy[k][m] - Axy[k-1][m]) * sinkingRate[vi] * vars[vi].current[km2ind(k,m)];
          }
          // partition it. There's loss across the bottom boundary, instantaneous recycling into another pool, and just sitting there
          Flux sedloss = getFlux("sedloss-" + vars[vi].shortname);
          Flux sedrecyc = getFlux("sedrecyc-" + vars[vi].shortname);
          for (int k=kbot[m]; k<K; k++) {
            sedloss.current[km2ind(k,m)] = F[k] * benthicLossFrac[vi] / V[k][m];
            sedrecyc.current[km2ind(k,m)] = F[k] * benthicRecyclingFrac[vi] / V[k][m];
          }          
        }
      }
    }    
  }


  // light attenuation ---------------------------------------------------------------- 
  
  
  float attTot(int k, int m) {
    float att = getParamVal("attSW");
    att += getParamVal("attP") * getVar("P").current[km2ind(k,m)];
    return att;
  }
  
  float surfacePAR() {
    return getParamVal("PAR0");
  }

  void calcLight() {
    Statevar PAR = getVar("PAR");
    if (PAR==null) return;
    for (int m=0; m<M; m++) {
      float running_par = surfacePAR();
      for (int k=K-1; k>= kbot[m]; k--) {
        float att = attTot(k,m);
        running_par *= exp(att * (z_rho[k] - z_w[k+1])); // attenuate down to cell center
        PAR.current[km2ind(k,m)] = running_par;
        running_par *= exp(att * (z_w[k] - z_rho[k])); // attenuate to bottom of cell
      }
    }
  }
  


}