import fs from 'fs';
import path from 'path';
import { fileURLToPath } from 'url';
import { parse } from 'csv-parse/sync';
import LogisticRegression from 'ml-logistic-regression';
import { Matrix } from 'ml-matrix';

const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

async function train() {
  console.log('Loading results.csv...');
  const csvPath = path.join(__dirname, 'results.csv');
  const fileContent = fs.readFileSync(csvPath, 'utf8');
  
  const records = parse(fileContent, {
    columns: true,
    skip_empty_lines: true
  });

  const X = [];
  const Y = [];

  for (const row of records) {
    if (row.type !== 'signal') continue; // Skip controls if they exist
    
    // Parse features
    const rsi = parseFloat(row.rsi);
    const atrPct = parseFloat(row.atrPct);
    const rvol = parseFloat(row.rvol);
    const macd = parseFloat(row.macd);
    const regime = parseFloat(row.regime_risk_on);
    
    // Skip rows with missing data
    if (isNaN(rsi) || isNaN(atrPct) || isNaN(rvol) || isNaN(macd) || isNaN(regime)) continue;

    X.push([
      1.0, // Intercept
      rsi / 100,
      atrPct / 10,
      rvol / 5,
      macd / 2,
      regime
    ]);

    // Target: 1 if win or win_partial, 0 otherwise
    const isWin = (row.outcome === 'win' || row.outcome === 'win_partial') ? 1 : 0;
    Y.push(isWin);
  }

  console.log(`Loaded ${X.length} valid training samples.`);

  const xMat = new Matrix(X);
  const yMat = Matrix.columnVector(Y);

  console.log('Training Logistic Regression model...');
  const logreg = new LogisticRegression({ numSteps: 1000, learningRate: 0.05 });
  logreg.train(xMat, yMat);

  const modelJson = logreg.toJSON();
  
  const outputPath = path.join(__dirname, 'model_weights.json');
  fs.writeFileSync(outputPath, JSON.stringify(modelJson, null, 2));
  
  console.log(`Model trained and saved to ${outputPath}`);
}

train().catch(console.error);
