import java.io.File;
import java.util.Arrays;
import java.util.Map;
import org.fireviewer.litert.training.OnDeviceLearning;

/** Real Android CPU exercise: two image sizes, train, durable save and optimizer resume. */
public final class FireDetectorRuntimeCheck {
  public static void main(String[] args) {
    File model=new File(args[0]), root=new File(args[1]);
    int threads=args.length>2?Integer.parseInt(args[2]):2;
    String schema="0000000000000000000000000000000000000000000000000000000000000000";
    Map<String,Object> square=TrainingRuntimeCheck.inputs("x",TrainingRuntimeCheck.image(128,128));
    Map<String,Object> rectangle=TrainingRuntimeCheck.inputs("x",TrainingRuntimeCheck.image(128,192));
    float[] target=new float[600];
    target[0]=.3f;target[1]=.3f;target[2]=.7f;target[3]=.7f;target[4]=0;target[5]=1;
    File checkpoint;float[] trained,continued;float firstLoss,secondLoss;
    System.out.println("FIRE_CHECK opening "+model.getName());
    try(OnDeviceLearning session=new OnDeviceLearning(model,schema,threads)) {
      System.out.println("FIRE_CHECK opened");
      float[] original=session.infer(square).get("detections").getValues();
      System.out.println("FIRE_CHECK infer "+original.length);
      firstLoss=session.train(square,target,.001f);
      trained=session.infer(square).get("detections").getValues();
      TrainingRuntimeCheck.check(!Arrays.equals(original,trained),"Training did not change outputs");
      checkpoint=session.save(root,"synthetic-android-runtime-fixture",1L);
      System.out.println("FIRE_CHECK saved loss="+firstLoss);
      secondLoss=session.train(rectangle,target,.001f);
      continued=session.infer(rectangle).get("detections").getValues();
      System.out.println("FIRE_CHECK continued loss="+secondLoss);
    }
    try(OnDeviceLearning resumed=new OnDeviceLearning(model,schema,threads)) {
      resumed.restore(checkpoint);
      TrainingRuntimeCheck.check(Arrays.equals(trained,resumed.infer(square).get("detections").getValues()),"Restored predictions differ");
      resumed.train(rectangle,target,.001f);
      TrainingRuntimeCheck.check(Arrays.equals(continued,resumed.infer(rectangle).get("detections").getValues()),"Resumed optimizer differs");
    }
    try(OnDeviceLearning invalid=new OnDeviceLearning(model,"1000000000000000000000000000000000000000000000000000000000000000",threads)) {
      boolean rejected=false;
      try { invalid.restore(checkpoint); } catch(IllegalArgumentException expected) { rejected=true; }
      TrainingRuntimeCheck.check(rejected,"Incompatible class schema accepted");
    }
    System.out.println("ANDROID_FIRE_TRAINING_PASS {\"runtime\":\""+org.tensorflow.lite.TensorFlowLite.runtimeVersion()+"\",\"mixed_shapes\":[[128,128],[128,192]],\"train\":true,\"restore_exact\":true,\"optimizer_resume_exact\":true,\"incompatible_schema_rejected\":true,\"first_loss\":"+firstLoss+",\"second_loss\":"+secondLoss+"}");
  }
}