import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.util.Locale;
public class LorenzLyapunov {



    private static double sigma = 10.0;
    private static double rho = 28.0;
    private static double beta = 8.0 / 3.0;


    private static final double DT = 0.001;
    private static final double T0 = 0.0;
    private static final double T_END = 50.0;



    private static final double FIT_T_MIN = 2.0;
    private static final double FIT_DELTA_STOP = 1.0;



    private static final double[] SEED = {1.0, 1.0, 1.0};
    private static double delta0X = 1e-5;


    private static final double DEFAULT_SPINUP = 30.0;

    public static void main(String[] args) throws IOException {
        Locale.setDefault(Locale.US);

        System.out.println("Java version: " + System.getProperty("java.version"));

        double spinUp = parseArguments(args);


        validateFixedPoints();


        double[] INIT_A = spinUp(SEED, spinUp);
        double[] INIT_B = {INIT_A[0] + delta0X, INIT_A[1], INIT_A[2]};

        System.out.println("Configuration:");
        System.out.printf("  sigma=%.4f  rho=%.4f  beta=%.6f%n", sigma, rho, beta);
        System.out.printf("  dt=%.4f  t_end=%.1f%n", DT, T_END);
        System.out.printf("  spin-up: %.3f time units from seed (%.1f, %.1f, %.1f)%n",
                spinUp, SEED[0], SEED[1], SEED[2]);
        System.out.printf("  state A after spin-up: (%.12f, %.12f, %.12f)%n",
                INIT_A[0], INIT_A[1], INIT_A[2]);
        System.out.printf("  perturbation: %.1e in x%n", delta0X);
        System.out.printf("  fit rule: first t > %.3f until first delta >= %.3f%n",
                FIT_T_MIN, FIT_DELTA_STOP);

        int steps = (int) Math.round((T_END - T0) / DT);


        double[] t = new double[steps + 1];

        double[] xA = new double[steps + 1];
        double[] yA = new double[steps + 1];
        double[] zA = new double[steps + 1];

        double[] xB = new double[steps + 1];
        double[] yB = new double[steps + 1];
        double[] zB = new double[steps + 1];

        double[] delta = new double[steps + 1];


        t[0] = T0;
        xA[0] = INIT_A[0]; yA[0] = INIT_A[1]; zA[0] = INIT_A[2];
        xB[0] = INIT_B[0]; yB[0] = INIT_B[1]; zB[0] = INIT_B[2];

        delta[0] = separation(xA[0], yA[0], zA[0], xB[0], yB[0], zB[0]);


        for (int i = 0; i < steps; i++) {
            t[i + 1] = t[i] + DT;

            double[] nextA = rk4Step(xA[i], yA[i], zA[i], DT);
            xA[i + 1] = nextA[0];
            yA[i + 1] = nextA[1];
            zA[i + 1] = nextA[2];

            double[] nextB = rk4Step(xB[i], yB[i], zB[i], DT);
            xB[i + 1] = nextB[0];
            yB[i + 1] = nextB[1];
            zB[i + 1] = nextB[2];

            delta[i + 1] = separation(xA[i + 1], yA[i + 1], zA[i + 1],
                    xB[i + 1], yB[i + 1], zB[i + 1]);
        }


        FitResult fit = fitLambdaLeastSquaresContiguous(t, delta, FIT_T_MIN, FIT_DELTA_STOP);

        double lambda = fit.slope;
        double doublingTime = Math.log(2.0) / lambda;
        double delta0 = delta[0];


        double tolerance = 1.0;
        double horizon = (1.0 / lambda) * Math.log(tolerance / delta0);

        System.out.println("Lyapunov fit results:");
        System.out.printf("  points used: %d%n", fit.n);
        System.out.printf("  fit t-range: [%.6f, %.6f]%n", fit.tMinUsed, fit.tMaxUsed);
        System.out.printf("  lambda (slope): %.8f%n", lambda);
        System.out.printf("  R^2: %.8f%n", fit.rSquared);
        System.out.printf("  intercept: %.8f%n", fit.intercept);
        System.out.printf("  doubling time ln(2)/lambda: %.8f%n", doublingTime);
        System.out.printf("  predictability horizon for Δ=%.3f: %.8f%n", tolerance, horizon);


        if (fit.tMaxUsed > 20.0) {
            System.out.println("WARNING: fit includes very late times; check window logic.");
        }


        String outFile = String.format(Locale.US, "lorenz_member_spinup%.0f.csv", spinUp);
        writeCsv(outFile, spinUp, INIT_A, t, xA, yA, zA, xB, yB, zB, delta);
        System.out.println("CSV written: " + outFile);
    }
    private static double parseArguments(String[] args) {
        if (args.length > 5) {
            throw new IllegalArgumentException(
                    "Usage: java LorenzLyapunov [spinUp] [sigma] [rho] [beta] [delta0]");
        }
        double spinUp = (args.length > 0) ? Double.parseDouble(args[0]) : DEFAULT_SPINUP;
        if (args.length > 1) sigma = Double.parseDouble(args[1]);
        if (args.length > 2) rho = Double.parseDouble(args[2]);
        if (args.length > 3) beta = Double.parseDouble(args[3]);
        if (args.length > 4) delta0X = Double.parseDouble(args[4]);
        if (spinUp < 0.0 || sigma <= 0.0 || rho <= 0.0 || beta <= 0.0 || delta0X <= 0.0) {
            throw new IllegalArgumentException("spinUp, sigma, rho, beta and delta0 must all be positive.");
        }
        return spinUp;
    }
    private static double[] spinUp(double[] seed, double tSpin) {
        double x = seed[0], y = seed[1], z = seed[2];
        int n = (int) Math.round(tSpin / DT);
        for (int i = 0; i < n; i++) {
            double[] next = rk4Step(x, y, z, DT);
            x = next[0]; y = next[1]; z = next[2];
        }
        return new double[]{x, y, z};
    }
    private static double[] deriv(double x, double y, double z) {
        double dx = sigma * (y - x);
        double dy = x * (rho - z) - y;
        double dz = x * y - beta * z;
        return new double[]{dx, dy, dz};
    }
    private static double[] rk4Step(double x, double y, double z, double dt) {
        double[] k1 = deriv(x, y, z);

        double[] k2 = deriv(
                x + 0.5 * dt * k1[0],
                y + 0.5 * dt * k1[1],
                z + 0.5 * dt * k1[2]
        );

        double[] k3 = deriv(
                x + 0.5 * dt * k2[0],
                y + 0.5 * dt * k2[1],
                z + 0.5 * dt * k2[2]
        );

        double[] k4 = deriv(
                x + dt * k3[0],
                y + dt * k3[1],
                z + dt * k3[2]
        );

        double xn = x + (dt / 6.0) * (k1[0] + 2.0 * k2[0] + 2.0 * k3[0] + k4[0]);
        double yn = y + (dt / 6.0) * (k1[1] + 2.0 * k2[1] + 2.0 * k3[1] + k4[1]);
        double zn = z + (dt / 6.0) * (k1[2] + 2.0 * k2[2] + 2.0 * k3[2] + k4[2]);

        return new double[]{xn, yn, zn};
    }
    private static double separation(double x1, double y1, double z1,
                                     double x2, double y2, double z2) {
        double dx = x1 - x2;
        double dy = y1 - y2;
        double dz = z1 - z2;
        return Math.sqrt(dx * dx + dy * dy + dz * dz);
    }
    private static FitResult fitLambdaLeastSquaresContiguous(double[] t, double[] delta,
                                                             double tMin, double deltaStop) {
        double sumX = 0.0, sumY = 0.0, sumXX = 0.0, sumXY = 0.0;
        int n = 0;

        boolean started = false;
        double tMinUsed = Double.NaN;
        double tMaxUsed = Double.NaN;

        for (int i = 0; i < t.length; i++) {
            if (!started) {
                if (t[i] > tMin) started = true;
                else continue;
            }

            if (delta[i] >= deltaStop) {
                break;
            }

            if (delta[i] <= 0.0) continue;

            double x = t[i];
            double y = Math.log(delta[i]);

            if (Double.isNaN(tMinUsed)) tMinUsed = x;
            tMaxUsed = x;

            sumX += x;
            sumY += y;
            sumXX += x * x;
            sumXY += x * y;
            n++;
        }

        if (n < 2) {
            throw new IllegalStateException("Not enough points in contiguous fit window. Adjust FIT_T_MIN / FIT_DELTA_STOP.");
        }

        double denom = n * sumXX - sumX * sumX;
        if (Math.abs(denom) < 1e-15) {
            throw new IllegalStateException("Degenerate least-squares denominator.");
        }

        double slope = (n * sumXY - sumX * sumY) / denom;
        double intercept = (sumY - slope * sumX) / n;



        double meanY = sumY / n;
        double ssRes = 0.0, ssTot = 0.0;
        boolean started2 = false;
        for (int i = 0; i < t.length; i++) {
            if (!started2) {
                if (t[i] > tMin) started2 = true;
                else continue;
            }
            if (delta[i] >= deltaStop) break;
            if (delta[i] <= 0.0) continue;
            double yObs = Math.log(delta[i]);
            double yPred = intercept + slope * t[i];
            ssRes += (yObs - yPred) * (yObs - yPred);
            ssTot += (yObs - meanY) * (yObs - meanY);
        }
        double rSquared = (ssTot > 0.0) ? 1.0 - ssRes / ssTot : Double.NaN;

        return new FitResult(slope, intercept, n, tMinUsed, tMaxUsed, rSquared);
    }
    private static void validateFixedPoints() {
        double s = Math.sqrt(beta * (rho - 1.0));
        double z = rho - 1.0;

        double[] fPlus = deriv(+s, +s, z);
        double[] fMinus = deriv(-s, -s, z);

        System.out.println("Fixed-point derivative checks (should be near 0):");
        System.out.printf("  f(+s,+s,z): [%.3e, %.3e, %.3e]%n", fPlus[0], fPlus[1], fPlus[2]);
        System.out.printf("  f(-s,-s,z): [%.3e, %.3e, %.3e]%n", fMinus[0], fMinus[1], fMinus[2]);
    }
    private static void writeCsv(String fileName, double spinUp, double[] stateA,
                                 double[] t,
                                 double[] xA, double[] yA, double[] zA,
                                 double[] xB, double[] yB, double[] zB,
                                 double[] delta) throws IOException {
        try (BufferedWriter bw = new BufferedWriter(new FileWriter(fileName))) {


            bw.write(String.format(Locale.US,
                    "# sigma=%.12g,rho=%.12g,beta=%.12g,dt=%.12g,t_end=%.12g,spin_up=%.12g,seed_x=%.12g,seed_y=%.12g,seed_z=%.12g,state_a_x=%.12g,state_a_y=%.12g,state_a_z=%.12g,delta0=%.12g,perturbation_axis=x,fit_t_min=%.12g,fit_delta_stop=%.12g%n",
                    sigma, rho, beta, DT, T_END, spinUp,
                    SEED[0], SEED[1], SEED[2], stateA[0], stateA[1], stateA[2],
                    delta0X, FIT_T_MIN, FIT_DELTA_STOP));
            bw.write("t,xA,yA,zA,xB,yB,zB,delta\n");
            for (int i = 0; i < t.length; i++) {
                bw.write(String.format(Locale.US,
                        "%.6f,%.12f,%.12f,%.12f,%.12f,%.12f,%.12f,%.12e%n",
                        t[i], xA[i], yA[i], zA[i], xB[i], yB[i], zB[i], delta[i]));
            }
        }
    }

    private static class FitResult {
        final double slope;
        final double intercept;
        final int n;
        final double tMinUsed;
        final double tMaxUsed;
        final double rSquared;

        FitResult(double slope, double intercept, int n, double tMinUsed, double tMaxUsed,
                  double rSquared) {
            this.slope = slope;
            this.intercept = intercept;
            this.n = n;
            this.tMinUsed = tMinUsed;
            this.tMaxUsed = tMaxUsed;
            this.rSquared = rSquared;
        }
    }
}