import Bun from 'bun';
import path from 'node:path';
import { parseArgs } from 'node:util';
import fs from 'node:fs';

const { values } = parseArgs({
    args: Bun.argv,
    options: {
        workingDir: { type: "string" },
        dataDir: { type: "string" },
        script: { type: "string" },
    },
    strict: true,
    allowPositionals: true,
});

if (!values.workingDir || !values.dataDir || !values.script) {
    console.error("Usage: bun run script/compare_factor.ts --workingDir <dir> --dataDir <dir> --script <python_script>");
    process.exit(1);
}

const workingDir = values.workingDir;
const dataDir = values.dataDir;
const predictScript = values.script;
const tempFileName = "temp.json";
const tempFilePath = path.join(workingDir, tempFileName);

// 1. 전처리 실행 (Factor 기반)
console.log("Step 1: Running factor preprocessing to temp.json...");
const preprocessResult = Bun.spawnSync([
    "bun", "run", "script/preprocess_factor.ts",
    "--workingDir", workingDir,
    "--dataDir", dataDir,
    "--fileName", tempFileName
]);

if (!preprocessResult.success) {
    console.error("Preprocessing failed");
    console.error(preprocessResult.stderr.toString());
    process.exit(1);
}

// 2. measure.csv 로드
console.log("Step 2: Loading measure.csv...");
const measurePath = path.join(dataDir, "measure.csv");
const measureContent = fs.readFileSync(measurePath, "utf-8");
const measureMap = new Map<string, number>();

measureContent.split("\n").forEach((line, index) => {
    if (index === 0 || !line.trim()) return;
    const parts = line.split(",");
    if (parts.length >= 3) {
        const constant = parts[0];
        const songno = parts[1];
        const diff = parts[2];
        measureMap.set(`${songno.trim()}_${diff.trim()}`, parseFloat(constant));
    }
});

// 3. temp.json 로드하여 대상 곡 목록 추출
const factors = JSON.parse(fs.readFileSync(tempFilePath, "utf-8"));
const uniqueSongnos = Array.from(new Set(factors.map((f: any) => f.songno)));

// 4. 예측 및 비교
console.log(`Step 3: Predicting and comparing ${uniqueSongnos.length} songs (Factor-based)...`);
const comparisonResults: any[] = [];
let processedCount = 0;

for (const songno of uniqueSongnos) {
    try {
        const predictProcess = Bun.spawnSync([
            "python3", predictScript,
            "--workingDir", workingDir,
            "--songno", songno as string,
            "--factor", tempFilePath // 파이썬 스크립트 인자명이 --factor로 고정되어 있는 경우를 가정
        ]);

        if (!predictProcess.success) {
            console.error(`\n[ERROR] Failed to predict songno ${songno}`);
            processedCount++;
            continue;
        }

        const output = predictProcess.stdout.toString().trim();
        const jsonStart = output.indexOf('[');
        const jsonEnd = output.lastIndexOf(']') + 1;

        if (jsonStart !== -1 && jsonEnd !== 0) {
            const predictions = JSON.parse(output.substring(jsonStart, jsonEnd));
            predictions.forEach((pred: any) => {
                const key = `${pred.songno}_${pred.diff}`;
                const actual = measureMap.get(key);
                if (actual !== undefined) {
                    comparisonResults.push({
                        songno: pred.songno,
                        diff: pred.diff,
                        actual: actual,
                        predicted: pred.predicted,
                        error: Math.abs(actual - pred.predicted)
                    });
                }
            });
        }

        processedCount++;
        if (processedCount % 10 === 0 || processedCount === uniqueSongnos.length) {
            const percent = ((processedCount / uniqueSongnos.length) * 100).toFixed(1);
            process.stdout.write(`\rProgress: ${processedCount}/${uniqueSongnos.length} (${percent}%)   `);
        }
    } catch (err) {
        console.error(`\nError processing songno ${songno}:`, err);
        processedCount++;
    }
}
console.log("\nPrediction finished.");

const avgError = comparisonResults.reduce((acc, curr) => acc + curr.error, 0) / comparisonResults.length;
const resultData = {
    summary: {
        total_compared: comparisonResults.length,
        average_absolute_error: avgError,
        timestamp: new Date().toISOString(),
        script_used: predictScript,
        type: "factor"
    },
    details: comparisonResults.sort((a, b) => b.error - a.error)
};

const comparePath = path.join(workingDir, "compare.json");
fs.writeFileSync(comparePath, JSON.stringify(resultData, null, 2), "utf-8");

console.log(`\nComparison complete! Results saved to: ${comparePath}`);

// 6. 결과 시각화 (compare.png 생성)
if (comparisonResults.length > 0) {
    console.log("Step 4: Generating comparison plot (compare.png)...");
    const plotPythonCode = `
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import json
import os

with open('${comparePath}', 'r', encoding='utf-8') as f:
    data = json.load(f)

if not data['details']:
    print("No details to plot.")
    exit(0)

df = pd.DataFrame(data['details'])
df['abs_error'] = (df['actual'] - df['predicted']).abs()
df = df.sort_values('abs_error', ascending=False).reset_index(drop=True)

plt.figure(figsize=(12, 6))
plt.switch_backend('Agg')
sns.scatterplot(x=df.index, y=df['abs_error'], alpha=0.6, s=20, color='royalblue')

plt.axhline(0.2, color='green', linestyle='--', linewidth=0.8, alpha=0.5, label='Target (0.2)')
plt.axhline(0.5, color='blue', linestyle='--', linewidth=0.8, alpha=0.5)
plt.axhline(1.0, color='red', linestyle='--', linewidth=0.8, alpha=0.5)

plt.ylim(0, max(3.5, df['abs_error'].max() + 0.5) if not df.empty else 3.5)
plt.title('Comparison Absolute Error Distribution (Feature-based)', fontsize=14)
plt.xlabel('Samples (Sorted by Error Magnitude)', fontsize=12)
plt.ylabel('Absolute Error', fontsize=12)
plt.grid(True, axis='y', alpha=0.3)

plot_path = os.path.join('${workingDir}', 'compare.png')
plt.savefig(plot_path)
print(f"Plot saved to: {plot_path}")
`;

    const plotProc = Bun.spawnSync(["python3", "-c", plotPythonCode]);
    if (plotProc.success) {
        console.log(plotProc.stdout.toString().trim());
    } else {
        console.error("Failed to generate plot:");
        console.error(plotProc.stderr.toString());
    }
}