#include "runtime.h"
#include <sundials/sundials_config.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

static void json_string(FILE *f, const char *value) {
    fputc('"',f);
    for (const unsigned char *p=(const unsigned char *)value; *p; p++) {
        if (*p=='"' || *p=='\\') { fputc('\\',f); fputc(*p,f); }
        else if (*p<32) fprintf(f,"\\u%04x",*p);
        else fputc(*p,f);
    }
    fputc('"',f);
}
static void vector(FILE *f, const double *value, int count) {
    fputc('[',f);
    for (int i=0;i<count;i++) fprintf(f,"%s%.17g",i?",":"",value[i]);
    fputc(']',f);
}
static int probe(void) {
    double t, y[NSTATES], dy[NSTATES], w[NOUTPUTS];
    while (scanf("%lf",&t)==1) {
        for (int i=0;i<NSTATES;i++) if (scanf("%lf",&y[i])!=1) return 64;
        int ok=model_eval(t,y,dy,w);
        printf("{\"success\":%s",ok?"true":"false");
        if (ok) { printf(",\"rhs\":"); vector(stdout,dy,NSTATES); printf(",\"outputs\":"); vector(stdout,w,NOUTPUTS); }
        printf("}\n");
    }
    return 0;
}
static int write_result(NativeRun *r, const char *path) {
    double dy[NSTATES], final[NOUTPUTS];
    int final_ok=model_eval(r->final_time,r->final_state,dy,final);
    size_t length=r->count*NOUTPUTS;
    if (length>268435456u/sizeof(double)) return 0;
    double *values=length?malloc(length*sizeof(double)):NULL;
    if (length && !values) return 0;
    for (size_t i=0;i<r->count;i++) {
        if (!model_eval(r->times[i],r->states+i*NSTATES,dy,values+i*NOUTPUTS)) {
            free(values); return 0;
        }
    }
    FILE *f=fopen(path,"wb"); if (!f) { free(values); return 0; }
    fprintf(f,"{\"success\":%s,\"status\":",!r->status?"true":"false");
    json_string(f,r->status==0?"completed":r->status==1?"cancelled":"failed");
    fprintf(f,",\"message\":"); json_string(f,r->message);
    fprintf(f,",\"backend\":\"native-c\",\"method\":\"%s\",\"solver\":\"%s\",\"sundialsVersion\":\"%s\","
              "\"simulatedUntil\":%.17g,\"solveSeconds\":%.17g,\"solveCpuSeconds\":%.17g,"
              "\"nfev\":%lu,\"acceptedSteps\":%lu,\"rejectedSteps\":%lu,\"stateTransitions\":%lu,"
              "\"solverStarts\":%lu,\"njev\":%lu,\"nlu\":%lu,\"maxAcceptedStep\":%.17g,\"series\":{",
            r->options.bdf?"BDF":"RK45",r->options.bdf?"CVODE":"Dormand-Prince 5(4)",SUNDIALS_VERSION,
            r->final_time,r->solve_seconds,r->solve_cpu_seconds,r->nfev,r->accepted,r->rejected,
            r->events,r->starts,r->njev,r->nlu,r->max_accepted_step);
    if (r->count) {
        fprintf(f,"\"time\":[");
        for (size_t i=0;i<r->count;i++) fprintf(f,"%s%.17g",i?",":"",r->times[i]);
        fputc(']',f);
        for (int j=0;j<NOUTPUTS;j++) {
            fputc(',',f); json_string(f,model_output_keys[j]); fprintf(f,":[");
            for (size_t i=0;i<r->count;i++) fprintf(f,"%s%.17g",i?",":"",values[i*NOUTPUTS+j]);
            fputc(']',f);
        }
    }
    fprintf(f,"},\"final\":{");
    if (final_ok) for (int j=0;j<NOUTPUTS;j++) {
        if (j) fputc(',',f);
        json_string(f,model_output_keys[j]); fprintf(f,":%.17g",final[j]);
    }
    fprintf(f,"},\"finalState\":"); vector(f,r->final_state,NSTATES);
    fprintf(f,"}\n"); int ok=!ferror(f); if (fclose(f)) ok=0;
    free(values); return ok;
}

int main(int argc, char **argv) {
    NativeRun r={0};
    r.options=(NativeOptions){0,10,.02,.001,1e-6,300,0,1,NULL};
    const char *output="result.json";
    for (int i=1;i<argc;i++) {
        const char *arg=argv[i];
        if (!strcmp(arg,"--probe")) return probe();
        if (!strcmp(arg,"--init")) {
            double y[NSTATES]; if (!model_init(y)) return 2;
            vector(stdout,y,NSTATES); fputc('\n',stdout); return 0;
        }
        if (!strcmp(arg,"--solve-only")) { r.options.record_samples=0; continue; }
        if (i+1==argc) return 64;
        const char *value=argv[++i];
        if (!strcmp(arg,"--output")) output=value;
        else if (!strcmp(arg,"--cancel-file")) r.options.cancel_path=value;
        else if (!strcmp(arg,"--method")) {
            if (strcmp(value,"RK45") && strcmp(value,"BDF")) return 64;
            r.options.bdf=!strcmp(value,"BDF");
        } else {
            char *end; double number=strtod(value,&end);
            if (*end || !isfinite(number)) return 64;
            if (!strcmp(arg,"--start")) r.options.start=number;
            else if (!strcmp(arg,"--stop")) r.options.stop=number;
            else if (!strcmp(arg,"--sample-step")) r.options.sample_step=number;
            else if (!strcmp(arg,"--max-step")) r.options.max_step=number;
            else if (!strcmp(arg,"--rtol")) r.options.rtol=number;
            else if (!strcmp(arg,"--timeout")) r.options.timeout=number;
            else return 64;
        }
    }
    if (!(r.options.stop>r.options.start && r.options.sample_step>0 && r.options.max_step>0 &&
          r.options.rtol>0 && r.options.timeout>0) ||
        r.options.start+r.options.sample_step==r.options.start ||
        r.options.start+r.options.max_step==r.options.start ||
        !isfinite(r.options.stop-r.options.start)) return 64;
    if (r.options.record_samples &&
        ((r.options.stop-r.options.start)/r.options.sample_step>1000000 ||
         ((r.options.stop-r.options.start)/r.options.sample_step+1024)*(NSTATES+NOUTPUTS)*sizeof(double)>268435456)) return 64;
    native_solve(&r);
    int saved=write_result(&r,output);
    int code=saved?(r.status==2?2:0):3;
    native_run_free(&r); return code;
}