package com.nablatensor.notebook;

import com.nablatensor.jax4j.tensor.*;
import com.nablatensor.jax4j.nn.FunctionalAdam;
import java.io.*;
import java.util.*;
import java.util.function.Function;
import java.util.zip.GZIPInputStream;

/** Dataset batching and typed visual output; the network and gradient step live in the cells. */
public final class Mnist {
  private Mnist() {}
  /** Eager Adam for the browser bridge, which cannot compile generated JIT kernels. */
  public static FunctionalAdam.Update<List<Tensor>> adam(List<Tensor> parameters,
      List<Tensor> gradients, FunctionalAdam.State<List<Tensor>> state, double rate) {
    List<Tensor> m = new ArrayList<>(), v = new ArrayList<>(), p = new ArrayList<>();
    Tensor step = state.step().add(1);
    try (Tensor b1scaled = step.mul(Math.log(.9)); Tensor b1exp = b1scaled.exp();
         Tensor b1neg = b1exp.neg(); Tensor bc1 = b1neg.add(1);
         Tensor b2scaled = step.mul(Math.log(.999)); Tensor b2exp = b2scaled.exp();
         Tensor b2neg = b2exp.neg(); Tensor bc2 = b2neg.add(1)) {
      for (int i = 0; i < parameters.size(); i++) {
        Tensor g = gradients.get(i);
        try (Tensor oldM = state.firstMoment().get(i).mul(.9); Tensor dg = g.mul(.1);
             Tensor oldV = state.secondMoment().get(i).mul(.999);
             Tensor gg = g.mul(g); Tensor dgg = gg.mul(.001)) {
          Tensor mm = oldM.add(dg), vv = oldV.add(dgg);
          m.add(mm); v.add(vv);
          try (Tensor mh = mm.div(bc1); Tensor vh = vv.div(bc2); Tensor root = vh.sqrt();
               Tensor denom = root.add(1e-8); Tensor direction = mh.div(denom);
               Tensor delta = direction.mul(rate)) {
            p.add(parameters.get(i).sub(delta));
          }
        }
      }
      return new FunctionalAdam.Update<>(List.copyOf(p),
          new FunctionalAdam.State<>(List.copyOf(m), List.copyOf(v), step));
    } catch (RuntimeException | Error failure) {
      m.forEach(Tensor::close); v.forEach(Tensor::close); p.forEach(Tensor::close);
      step.close(); throw failure;
    }
  }
  public record Batch(Tensor x, Tensor y) implements AutoCloseable {
    public void close() { x.close(); y.close(); }
  }
  @FunctionalInterface public interface Step { double apply(Tensor x, Tensor y); }
  public static final class Data {
    public final int count;
    private final byte[] rows;
    private Data(int count, byte[] rows) { this.count=count; this.rows=rows; }
    public int label(int i) { return rows[i*785]&255; }
    public byte[] pixels(int i) { return Arrays.copyOfRange(rows,i*785+1,(i+1)*785); }
    public int[] classCounts() {
      int[] counts=new int[10]; for(int i=0;i<count;i++)counts[label(i)]++;
      return counts;
    }
    /** Deterministic equal-per-class subset, retaining source order and original pixels. */
    public Data balanced(int selectedCount) {
      if(selectedCount<10 || selectedCount>count || selectedCount%10!=0)
        throw new IllegalArgumentException("Balanced sample count must be a multiple of 10 between 10 and "+count);
      int quota=selectedCount/10;
      for(int available:classCounts())if(available<quota)
        throw new IllegalArgumentException("Not enough images in every digit class");
      if(selectedCount==count)return this;
      byte[] selected=new byte[selectedCount*785];int[] used=new int[10];int at=0;
      for(int i=0;i<count && at<selectedCount;i++) {
        int digit=label(i);
        if(used[digit]<quota){System.arraycopy(rows,i*785,selected,at++*785,785);used[digit]++;}
      }
      return new Data(selectedCount,selected);
    }
    public Batch batch(int[] order, int offset, int n, int pool) {
      if (pool!=1 && pool!=2 && pool!=4) throw new IllegalArgumentException("pool must be 1, 2 or 4");
      int side=28/pool, features=side*side;
      float[] x=new float[n*features], y=new float[n*10];
      for (int b=0;b<n;b++) {
        int row=order[offset+b];
        if(row<0 || row>=count) throw new IllegalArgumentException("Sample index out of bounds");
        y[b*10+label(row)]=1;
        for(int r=0;r<side;r++) for(int c=0;c<side;c++) {
          int sum=0;
          for(int dr=0;dr<pool;dr++) for(int dc=0;dc<pool;dc++)
            sum+=rows[row*785+1+(r*pool+dr)*28+c*pool+dc]&255;
          x[b*features+r*side+c]=(float)sum/(255*pool*pool);
        }
      }
      return new Batch(Tensors.array(x,n,features),Tensors.array(y,n,10));
    }
    public double epoch(int count, int batchSize, int pool, long seed, Step step) {
      if(count<1 || count>this.count || batchSize<1 || batchSize>256) throw new IllegalArgumentException("Invalid training count or batch size");
      Data selected=balanced(count);
      int[] order=new int[count]; for(int i=0;i<count;i++) order[i]=i;
      Random random=new Random(seed);
      for(int i=count-1;i>0;i--) {int j=random.nextInt(i+1),v=order[i];order[i]=order[j];order[j]=v;}
      double total=0;
      for(int at=0;at<count;at+=batchSize) {
        int n=Math.min(batchSize,count-at);
        try(Batch b=selected.batch(order,at,n,pool)) {total+=step.apply(b.x(),b.y())*n;}
      }
      return total/count;
    }
    public double loss(int count,int pool,Step loss) {
      if(count<1 || count>this.count) throw new IllegalArgumentException("Invalid evaluation count");
      Data selected=balanced(count);
      int[] order=new int[count];for(int i=0;i<count;i++)order[i]=i;
      double total=0;
      for(int at=0;at<count;at+=64) {
        int n=Math.min(64,count-at);
        try(Batch b=selected.batch(order,at,n,pool)) {total+=loss.apply(b.x(),b.y())*n;}
      }
      return total/count;
    }
    public Report evaluate(int pool,Function<Tensor,Tensor> predict) {
      int[][] confusion=new int[10][10];int correct=0;
      balanced(count); // Evaluation must cover an equal number of every digit.
      int[] examples=new int[10], predicted=new int[10];boolean[] hasError=new boolean[10];
      Arrays.fill(examples,-1);
      int[] order=new int[count];for(int i=0;i<count;i++)order[i]=i;
      for(int at=0;at<count;at+=64) {
        int n=Math.min(64,count-at);
        try(Batch b=batch(order,at,n,pool);Tensor output=predict.apply(b.x())) {
          float[] values=output.toFloatArray();
          for(int i=0;i<n;i++) {
            int best=0;for(int j=1;j<10;j++)if(values[i*10+j]>values[i*10+best])best=j;
            int actual=label(at+i);confusion[actual][best]++;
            if(best==actual)correct++;
            if(examples[actual]<0 || (!hasError[actual] && best!=actual)) {
              examples[actual]=at+i;predicted[actual]=best;
              if(best!=actual)hasError[actual]=true;
            }
          }
        }
      }
      List<Integer> samples=new ArrayList<>(), guesses=new ArrayList<>();
      for(int digit=0;digit<10;digit++){samples.add(examples[digit]);guesses.add(predicted[digit]);}
      return new Report(this,correct,confusion,samples,guesses);
    }
    public void showSamples() {
      StringBuilder out=new StringBuilder("{\"kind\":\"samples\",\"items\":[");
      Data samples=balanced(10);
      for(int i=0;i<10;i++){if(i>0)out.append(',');image(out,samples,i,samples.label(i));}
      emit(out.append("]}").toString());
    }
  }
  public record Report(Data data,int correct,int[][] confusion,List<Integer> samples,List<Integer> guesses) {
    public void show() {
      StringBuilder out=new StringBuilder("{\"kind\":\"evaluation\",\"correct\":").append(correct)
        .append(",\"total\":").append(data.count).append(",\"confusion\":").append(Arrays.deepToString(confusion))
        .append(",\"items\":[");
      for(int i=0;i<samples.size();i++){if(i>0)out.append(',');image(out,data,samples.get(i),guesses.get(i));}
      emit(out.append("]}").toString());
    }
  }
  private static void image(StringBuilder out,Data data,int id,int guess) {
    out.append("{\"actual\":").append(data.label(id)).append(",\"predicted\":").append(guess)
      .append(",\"pixels\":\"").append(Base64.getEncoder().encodeToString(data.pixels(id))).append("\"}");
  }
  public static Data load(String split) throws IOException {
    if(!Set.of("train","validation","test").contains(split))throw new IllegalArgumentException("Unknown MNIST split");
    InputStream resource=Mnist.class.getResourceAsStream("/mnist/"+split+".bin.gz");
    if(resource==null)throw new FileNotFoundException("Missing MNIST resource: "+split);
    try(DataInputStream in=new DataInputStream(new GZIPInputStream(resource))) {
      if(in.readInt()!=0x4d4e5354)throw new IOException("Bad MNIST header");
      int count=in.readInt();if(count<1||count>10000)throw new IOException("Bad sample count");
      byte[] rows=in.readNBytes(count*785);
      if(rows.length!=count*785 || in.read()!=-1)throw new IOException("Bad MNIST length");
      for(int i=0;i<count;i++)if((rows[i*785]&255)>9)throw new IOException("Bad MNIST label");
      return new Data(count,rows);
    }
  }
  /** Independent double-precision reference reduction over the model's F32 logits. */
  public static double crossEntropy(Tensor logits, Tensor labels) {
    float[] z=logits.toFloatArray(), y=labels.toFloatArray();
    int n=logits.shape().dim(0);
    double total=0;
    for(int i=0;i<n;i++) {
      double max=Double.NEGATIVE_INFINITY;
      for(int j=0;j<10;j++)max=Math.max(max,z[i*10+j]);
      double sum=0;for(int j=0;j<10;j++)sum+=Math.exp(z[i*10+j]-max);
      double logSum=Math.log(sum);
      for(int j=0;j<10;j++)total-=y[i*10+j]*(z[i*10+j]-max-logSum);
    }
    return total/n;
  }
  public static Tensor weights(Random random,int in,int out) {
    float[] w=new float[in*out];double scale=Math.sqrt(2.0/in);
    for(int i=0;i<w.length;i++)w[i]=(float)(random.nextGaussian()*scale);
    return Tensors.array(w,in,out);
  }
  public static void curve(List<Double> training,List<Double> validation) {
    emit("{\"kind\":\"curve\",\"training\":"+training+",\"validation\":"+validation+"}");
  }
  private static void emit(String json) {System.out.println("@@MNIST@@"+json);}
}