import { parseTja } from './parse';
import { factorize, Factors } from './factorize';
import { join } from 'path';
import { readFileSync, writeFileSync, readdirSync } from 'fs';
import { parse } from 'csv-parse/sync';

async function train() {
    const args = process.argv.slice(2);
    const workingDir = args[0];
    const scriptDir = args[1];
    const dataDir = args[2];
    const trainCount = parseInt(args[3]) || 100;
    const validateCount = parseInt(args[4]) || 20;
    const margin = parseFloat(args[5]) || 0.1;

    const factorJsonPath = join(scriptDir, 'factor.json');
 // 항상 script/ 위치 참조
    let weights: Factors = JSON.parse(readFileSync(factorJsonPath, 'utf-8'));
    
    // 결과물(로그 등)을 저장할 경로 (필요 시 활용)
    const logPath = join(workingDir, 'training_log.txt'); 

    const measureCsv = readFileSync(join(dataDir, 'measure.csv'), 'utf-8');
    const records = parse(measureCsv, { columns: true, skip_empty_lines: true });

    const tjaFiles = readdirSync(join(dataDir, 'tja')).filter(f => f.endsWith('.tja'));
    const shuffled = tjaFiles.sort(() => 0.5 - Math.random());
    const trainFiles = shuffled.slice(0, trainCount);
    const validateFiles = shuffled.slice(trainCount, trainCount + validateCount);

    console.log(`Training with ${trainFiles.length} files...`);

    // 학습 로직 (Sigmoid + [1, 12] scaling)
    const sigmoid = (x: number) => 1 / (1 + Math.exp(-x));
    
    let error = Infinity;
    let iterations = 0;
    while (error > margin && iterations < 10) { // Spec에 따라 10번 반복
        let totalError = 0;
        let count = 0;

        for (const file of trainFiles) {
            const songno = file.replace(/\D/g, '');
            const tjaContent = readFileSync(join(dataDir, 'tja', file), 'utf-8');
            const parsed = parseTja(tjaContent);
            if (!parsed) continue;

            for (const diff of ['oni', 'edit'] as const) {
                const course = parsed[diff];
                if (!course) continue;

                const target = records.find((r: any) => r.songno === songno && (r.diff === (diff === 'oni' ? 'oni' : 'ura')));
                if (!target) continue;

                const factors = factorize(course);
                // 가중치 적용 전 정규화 (임의 값으로 가정)
                const normalizedFactors = {
                    physical_density: Math.min(factors.physical_density / 20, 1),
                    stamina_requirement: Math.min(factors.stamina_requirement, 1),
                    pattern_complexity: Math.min(factors.pattern_complexity, 1),
                    rhythmic_complexity: Math.min(factors.rhythmic_complexity, 1),
                    reading_gimmick: Math.min(factors.reading_gimmick, 1)
                };

                const rawPrediction = Object.keys(normalizedFactors).reduce((sum, key) => 
                    sum + (normalizedFactors[key as keyof Factors] * weights[key as keyof Factors]), 0);
                
                const prediction = (sigmoid(rawPrediction) * 11) + 1; // [1, 12]
                
                const targetValue = parseFloat(target.상수);
                const diff_val = targetValue - prediction;
                totalError += Math.abs(diff_val);
                count++;

                // 가중치 업데이트 (간단한 경사 하강법)
                for (const key in weights) {
                    const k = key as keyof Factors;
                    weights[k] += diff_val * normalizedFactors[k] * 0.01;
                }
            }
        }
        error = totalError / (count || 1);
        console.log(`Iteration ${iterations}: Mean Error = ${error.toFixed(4)}`);
        iterations++;
    }

    writeFileSync(factorJsonPath, JSON.stringify(weights, null, 2));
    writeFileSync(join(workingDir, 'training_result.json'), JSON.stringify({ finalError: error, weights }, null, 2));
    console.log(`Training complete. Weights saved to ${factorJsonPath}, result saved to ${workingDir}`);

    // 검증 로직
    console.log('\nValidation Results:');
    for (const file of validateFiles) {
        const songno = file.replace(/\D/g, '');
        const tjaContent = readFileSync(join(dataDir, 'tja', file), 'utf-8');
        const parsed = parseTja(tjaContent);
        if (!parsed) continue;

        for (const diff of ['oni', 'edit'] as const) {
            const course = parsed[diff];
            if (!course) continue;

            const target = records.find((r: any) => r.songno === songno && (r.diff === (diff === 'oni' ? 'oni' : 'ura')));
            if (!target) {
                console.log(`[!] No match for ${songno} diff=${diff === 'oni' ? 'oni' : 'ura'}`);
                continue;
            }

            const factors = factorize(course);
            const prediction = Object.keys(factors).reduce((sum, key) => 
                sum + (factors[key as keyof Factors] * weights[key as keyof Factors]), 0);
            
            console.log(`[${songno}] Target: ${target.상수}, Predicted: ${prediction.toFixed(2)}, Diff: ${Math.abs(parseFloat(target.상수) - prediction).toFixed(2)}`);
        }
    }
}

train();