import { parseTja } from './parse';
import { factorize } from './factorize';
import { join } from 'path';
import { readFileSync, writeFileSync, unlinkSync, existsSync } from 'fs';
import { execSync } from 'child_process';

/**
 * 예측 엔진: 모델 실행을 위해 Python 환경을 호출
 */
function runInference(workingDir: string, factors: Record<string, number>): number {
    const tempFactorPath = join(workingDir, `temp_feat_${Date.now()}.json`);
    writeFileSync(tempFactorPath, JSON.stringify(factors));

    const pythonScript = `
import json, torch, sys
import numpy as np
import torch.nn as nn

# 모델 정의 (학습된 구조와 동일해야 함)
class DifficultyNet(nn.Module):
    def __init__(self, input_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Dropout(0.4),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()
        )
    def forward(self, x): return self.net(x)

def predict():
    try:
        with open('${tempFactorPath}', 'r') as f: factors = json.load(f)
        with open('${join(workingDir, 'scaler.json')}', 'r') as f: s = json.load(f)

        X = np.array([[factors[col] for col in s['cols']]], dtype=np.float32)
        
        # 정규화 (Robust/Standard 지원)
        if 'center' in s: X_scaled = (X - np.array(s['center'])) / np.array(s['scale'])
        else: X_scaled = (X - np.array(s.get('mean', 0))) / np.array(s.get('std', s.get('scale', 1)))

        model = DifficultyNet(len(s['cols']))
        model.load_state_dict(torch.load('${join(workingDir, 'model.pth')}', map_location='cpu'))
        model.eval()

        with torch.no_grad():
            input_tensor = torch.from_numpy(X_scaled).float()
            pred = model(input_tensor) * 11 + 1
        print(float(pred.item()))
    except Exception as e:
        print(f"ERROR: {e}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__": predict()
`;

    const tempPyPath = join(workingDir, `inference_${Date.now()}.py`);
    writeFileSync(tempPyPath, pythonScript);

    try {
        const result = execSync(`python3 ${tempPyPath}`).toString().trim();
        return parseFloat(result);
    } finally {
        if (existsSync(tempFactorPath)) unlinkSync(tempFactorPath);
        if (existsSync(tempPyPath)) unlinkSync(tempPyPath);
    }
}

function main() {
    const [workingDir, tjaPath, difficulty = 'oni'] = process.argv.slice(2);
    
    if (!workingDir || !tjaPath) {
        console.log("Usage: bun run script/predict.ts <workingDir> <tjaPath> [difficulty]");
        process.exit(1);
    }

    const tjaContent = readFileSync(tjaPath, 'utf-8');
    const parsed = parseTja(tjaContent);
    const chart = parsed[difficulty.toLowerCase() as 'oni' | 'edit'];

    if (!chart) {
        console.error(`Error: Difficulty '${difficulty}' not found.`);
        process.exit(1);
    }

    const score = runInference(workingDir, factorize(chart));
    console.log(`\n🎯 Predicted Difficulty: ${score.toFixed(2)}`);
}

main();