import pandas as pd
import json, os, sys
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.preprocessing import StandardScaler

MAX_MARGIN_LIMIT = 3

class DifficultyNet(nn.Module):
    def __init__(self, input_dim):
        super(DifficultyNet, self).__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 train(script_dir, working_dir, data_dir, train_count, val_count, margin, val_iterations):
    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
    print(f"🚀 Using Device: {device}")

    # 데이터 로드 및 병합
    factors_path = os.path.join(working_dir, 'factors.json')
    if not os.path.exists(factors_path):
        os.system(f"bun run {os.path.join(script_dir, 'factorize.ts')} {data_dir} {working_dir}")

    df_feat = pd.read_json(factors_path)
    df_meas = pd.read_csv(os.path.join(data_dir, 'measure.csv'))
    df_meas['diff'] = df_meas['diff'].replace('ura', 'edit')
    df_feat['songno'], df_meas['songno'] = df_feat['songno'].astype(int), df_meas['songno'].astype(int)
    df = pd.merge(df_feat, df_meas, on=['songno', 'diff'])

    exclude_cols = ['songno', 'diff', 'title', 'course', '상수', 'predicted_measure', 'error']
    X_cols = [c for c in df_feat.columns if c not in exclude_cols]
    
    scaler = StandardScaler()
    df[X_cols] = scaler.fit_transform(df[X_cols])
    with open(os.path.join(working_dir, 'scaler.json'), 'w') as f:
        json.dump({'mean': scaler.mean_.tolist(), 'std': scaler.scale_.tolist(), 'cols': X_cols}, f)

    model_path = os.path.join(working_dir, 'model.pth')
    margin = float(margin)
    val_iterations = int(val_iterations)
    
    model = DifficultyNet(len(X_cols)).to(device)
    attempt = 1

    while True:
        print(f"\n[Attempt {attempt}] " + "-"*40)
        if os.path.exists(model_path):
            try: model.load_state_dict(torch.load(model_path, map_location=device))
            except: pass

        train_df = df.sample(n=min(int(train_count), len(df)))
        X_train = torch.FloatTensor(train_df[X_cols].values).to(device)
        y_train = torch.FloatTensor(((train_df['상수'].values - 1) / 11)).unsqueeze(1).to(device)

        base_lr = 0.001
        optimizer = optim.Adam(model.parameters(), lr=base_lr, weight_decay=1e-5)
        criterion = nn.L1Loss()

        model.train()
        epoch = 0
        while True:
            optimizer.zero_grad()
            preds_train = model(X_train)
            loss = criterion(preds_train, y_train) # MSELoss 권장
            loss.backward()
            optimizer.step()
            
            # 현재 학습 데이터에 대한 오차 분석
            with torch.no_grad():
                diff = torch.abs((preds_train * 11 + 1) - (y_train * 11 + 1))
                train_mae = torch.mean(diff).item()
                train_max = torch.max(diff).item()
            
            if epoch % 1000 == 0:
                print(f"   - Ep {epoch:5d} | MAE: {train_mae:.4f} | MAX: {train_max:.4f}")
            
            # ✅ 조건 강화: MAE뿐만 아니라 MAX도 어느 정도 잡혔을 때만 검증으로 이동
            if train_mae < margin and train_max < (margin * MAX_MARGIN_LIMIT): 
                print(f"   ✅ Train Goal Reached (MAE: {train_mae} < {margin:.4f}, MAX: {train_max:.4f} < {(margin * MAX_MARGIN_LIMIT):.4f}). Moving to Val.")
                break
            
            epoch += 1
            if epoch > 50000:
                print("   ⚠️ Timed out. Resampling...")
                break
        
        if train_mae >= margin:
            torch.save(model.state_dict(), model_path)
            attempt += 1
            continue

        # 3. 검증 단계
        print(f"3. Validating {val_iterations} iterations...")
        model.eval()
        all_passed = True
        with torch.no_grad():
            for i in range(1, val_iterations + 1):
                val_df = df.sample(n=min(int(val_count), len(df)))
                X_val = torch.FloatTensor(val_df[X_cols].values).to(device)
                y_val_raw = val_df['상수'].values
                preds = model(X_val) * 11 + 1
                y_val_tensor = torch.FloatTensor(y_val_raw).unsqueeze(1).to(device)
                diff_tensor = torch.abs(preds - y_val_tensor)
                mae = torch.mean(diff_tensor).item()
                max_error = torch.max(diff_tensor).item()
                

                # csv 저장
                # preds를 CPU 넘파이로 변환하고 1차원으로 펴주는 과정이 포함되어야 합니다.
                preds_np = preds.detach().cpu().numpy().flatten() if torch.is_tensor(preds) else preds
                
                res_df = val_df[['songno', 'diff', '상수']].copy()
                res_df['predicted_measure'] = preds_np
                res_df['error'] = np.abs(preds_np - y_val_raw)
                
                output_file = os.path.join(working_dir, f'validate_result_{i}.csv')
                res_df.to_csv(output_file, index=False)

                if mae <= margin and max_error <= (margin * MAX_MARGIN_LIMIT):
                    print(f"   [Iter {i}] ✅ PASS (MAE: {mae:.4f}, MAX: {max_error:.4f})")
                else:
                    print(f"   [Iter {i}] ❌ FAIL (MAE: {mae:.4f} > {margin}, MAX: {max_error:.4f} > {margin * MAX_MARGIN_LIMIT})")
                    all_passed = False
                    break 

        if all_passed:
            print(f"\n✨ Final Success!")
            torch.save(model.state_dict(), model_path)
            return
        else:
            torch.save(model.state_dict(), model_path)
            attempt += 1

if __name__ == "__main__":
    if len(sys.argv) < 8: sys.exit(1)
    train(*sys.argv[1:8])