import fs from 'fs';

function sigmoid(z) {
  // Cap z to avoid overflow
  if (z > 20) return 1.0;
  if (z < -20) return 0.0;
  return 1 / (1 + Math.exp(-z));
}

function computeLogLoss(features, labels, weights) {
  let loss = 0;
  for (let i = 0; i < features.length; i++) {
    const f = features[i];
    let z = weights[0];
    for (let j = 0; j < f.length; j++) z += weights[j + 1] * f[j];
    const p = sigmoid(z);
    // clip p to avoid log(0)
    const pSafe = Math.max(1e-15, Math.min(1 - 1e-15, p));
    loss += -labels[i] * Math.log(pSafe) - (1 - labels[i]) * Math.log(1 - pSafe);
  }
  return loss / features.length;
}

export class LogisticModel {
  constructor() {
    this.weights = null;
    this.featureKeys = [];
  }

  // Internal full-batch gradient descent
  _train(X, Y, lambda, learningRate, maxSteps, tol) {
    const numFeatures = X[0].length;
    let w = new Array(numFeatures + 1).fill(0.0);
    const N = X.length;

    let prevLoss = Infinity;

    for (let step = 0; step < maxSteps; step++) {
      const grad = new Array(numFeatures + 1).fill(0.0);
      
      for (let i = 0; i < N; i++) {
        let z = w[0];
        for (let j = 0; j < numFeatures; j++) z += w[j + 1] * X[i][j];
        
        const p = sigmoid(z);
        const err = p - Y[i];
        
        grad[0] += err;
        for (let j = 0; j < numFeatures; j++) {
          grad[j + 1] += err * X[i][j];
        }
      }

      // Add L2 penalty (skip intercept grad[0])
      for (let j = 0; j < numFeatures; j++) {
        grad[j + 1] += lambda * w[j + 1];
      }

      // Update weights
      let maxDelta = 0;
      for (let j = 0; j < w.length; j++) {
        const update = (learningRate * grad[j]) / N;
        w[j] -= update;
        if (Math.abs(update) > maxDelta) maxDelta = Math.abs(update);
      }

      if (maxDelta < tol) {
        // Converged
        break;
      }
    }
    return w;
  }

  fit(featuresArray, labels) {
    if (!featuresArray.length) return;
    this.featureKeys = Object.keys(featuresArray[0]);
    
    // Convert to 2D array
    const X = featuresArray.map(f => this.featureKeys.map(k => f[k]));
    const Y = labels;
    const N = X.length;

    // 1. Carve a 80/20 validation slice from TRAIN to tune lambda
    const splitIdx = Math.floor(N * 0.8);
    const X_train = X.slice(0, splitIdx);
    const Y_train = Y.slice(0, splitIdx);
    const X_val = X.slice(splitIdx);
    const Y_val = Y.slice(splitIdx);

    // 2. Tune lambda
    const lambdas = [0.01, 0.1, 1.0, 10.0, 100.0, 500.0, 1000.0];
    let bestLambda = 0;
    let bestLoss = Infinity;
    
    console.log(`\nTuning L2 Regularization (Lambda) on Validation Slice...`);
    for (const l of lambdas) {
      const w = this._train(X_train, Y_train, l, 0.5, 2000, 1e-5);
      const loss = computeLogLoss(X_val, Y_val, w);
      console.log(`  Lambda ${l.toString().padEnd(6)} -> LogLoss: ${loss.toFixed(4)}`);
      if (loss < bestLoss) {
        bestLoss = loss;
        bestLambda = l;
      }
    }
    
    console.log(`Selected Best Lambda: ${bestLambda}`);

    // 3. Refit on FULL TRAIN set using best lambda
    this.weights = this._train(X, Y, bestLambda, 0.5, 5000, 1e-5);

    return this.getCoefficients();
  }

  predict(featuresArray) {
    if (!this.weights) throw new Error("Model not trained");
    
    const probs = [];
    for (let i = 0; i < featuresArray.length; i++) {
      const f = featuresArray[i];
      let z = this.weights[0]; // intercept
      for (let j = 0; j < this.featureKeys.length; j++) {
        z += this.weights[j + 1] * f[this.featureKeys[j]];
      }
      probs.push(sigmoid(z));
    }
    
    return probs;
  }

  getCoefficients() {
    if (!this.weights) return {};
    const coeffs = { intercept: this.weights[0] };
    for (let i = 0; i < this.featureKeys.length; i++) {
      coeffs[this.featureKeys[i]] = this.weights[i + 1];
    }
    return coeffs;
  }

  save(filepath) {
    fs.writeFileSync(filepath, JSON.stringify({
      featureKeys: this.featureKeys,
      weights: this.weights
    }, null, 2));
  }

  load(filepath) {
    if (!fs.existsSync(filepath)) return;
    const data = JSON.parse(fs.readFileSync(filepath, 'utf8'));
    this.featureKeys = data.featureKeys;
    this.weights = data.weights;
  }
}