#!/usr/bin/env python3
"""
Exp103: DINOv2 Giant + Ridge Grid Search

Ridgeのハイパーパラメータをグリッドサーチで最適化:
- alpha: 正則化強度 [0.1, 0.5, 1.0, 5.0, 10.0]
- solver: ソルバー ['auto', 'svd', 'cholesky']
- tol: 収束判定閾値 [1e-4, 1e-3]

Ridge回帰のL2正則化とソルバーの組み合わせで最適化。

期待OOF: 0.69-0.70
"""

import os
import gc
import numpy as np
import pandas as pd
from tqdm import tqdm
import torch
from transformers import AutoImageProcessor, AutoModel
from sklearn.linear_model import Ridge
from sklearn.model_selection import KFold
from PIL import Image
from itertools import product
import warnings

warnings.filterwarnings('ignore')
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
ROOT = "/kaggle/input/csiro-biomass/"
LABELS = ['Dry_Clover_g', 'Dry_Dead_g', 'Dry_Green_g', 'Dry_Total_g', 'GDM_g']
WEIGHTS = {'Dry_Green_g': 0.1, 'Dry_Dead_g': 0.1, 'Dry_Clover_g': 0.1, 'GDM_g': 0.2, 'Dry_Total_g': 0.5}

# Grid search parameters
PARAM_GRID = {
    'alpha': [0.1, 0.5, 1.0, 5.0, 10.0],
    'solver': ['auto', 'svd', 'cholesky'],
    'tol': [1e-4, 1e-3]
}

print("="*70)
print("🚀 Exp103: DINOv2 + Ridge Grid Search")
print("="*70)
print(f"Device: {DEVICE}")
print("Baseline: Exp101 Ridge alpha=0.5")
print("\nGrid Search Parameters:")
print(f"  alpha:  {PARAM_GRID['alpha']}")
print(f"  solver: {PARAM_GRID['solver']}")
print(f"  tol:    {PARAM_GRID['tol']}")
print(f"  Total combinations: {len(list(product(*PARAM_GRID.values())))}")
print("="*70)

# Load DINOv2 model
print("\n📦 Loading DINOv2-Giant model...")
processor = AutoImageProcessor.from_pretrained('facebook/dinov2-giant')
model = AutoModel.from_pretrained('facebook/dinov2-giant')
model = model.to(DEVICE)
print("✅ Model loaded")

# Load training data
print("\n📂 Loading training data...")
train_df = pd.read_csv(os.path.join(ROOT, "train.csv"))
print(f"Training samples: {len(train_df)}")

# Extract features
print("\n🔍 Extracting DINOv2 features from training images...")
embeds, targets = [], [[] for _ in range(5)]
target_mapping = {label: i for i, label in enumerate(LABELS)}

for img_path in tqdm(train_df['image_path'].unique(), desc="  Processing images"):
    full_path = os.path.join(ROOT, img_path)

    with Image.open(full_path) as img:
        inputs = processor(images=img, return_tensors="pt").to(DEVICE)

    with torch.no_grad():
        outputs = model(**inputs)
        img_name = os.path.splitext(os.path.basename(img_path))[0]
        embed = outputs.pooler_output.cpu()
        embeds.append(embed)

    # Get targets for this image
    img_rows = train_df[train_df['image_path'] == img_path]
    for _, row in img_rows.iterrows():
        target_idx = target_mapping[row['target']]
        targets[target_idx].append(torch.tensor([row['value']]))

print(f"✅ Extracted {len(embeds)} image embeddings (1536-dim)")

# Prepare data
embeds_np = np.array(torch.cat(embeds))

# Grid Search
print("\n" + "="*70)
print("🔍 Grid Search for Best Ridge Parameters")
print("="*70)

best_params = {}
best_scores = {}
all_results = []

for i, label in enumerate(LABELS):
    print(f"\n📊 Target: {label}")
    targets_np = np.array(torch.cat(targets[i]))

    best_score = -np.inf
    best_param = None

    # Test all parameter combinations
    for params in product(*PARAM_GRID.values()):
        alpha, solver, tol = params

        # 5-fold CV
        kf = KFold(n_splits=5, shuffle=True, random_state=42)
        fold_scores = []

        for fold, (train_idxs, val_idxs) in enumerate(kf.split(embeds_np)):
            X_train, X_val = embeds_np[train_idxs], embeds_np[val_idxs]
            y_train, y_val = targets_np[train_idxs], targets_np[val_idxs]

            # Train Ridge
            reg = Ridge(alpha=alpha, solver=solver, tol=tol, max_iter=10000)
            reg.fit(X_train, y_train)

            # R² on validation
            preds = reg.predict(X_val)
            ss_res = np.sum((y_val - preds) ** 2)
            ss_tot = np.sum((y_val - np.mean(y_val)) ** 2)
            r2 = 1 - (ss_res / ss_tot)
            fold_scores.append(r2)

        avg_score = np.mean(fold_scores)
        all_results.append({
            'target': label,
            'alpha': alpha,
            'solver': solver,
            'tol': tol,
            'r2': avg_score
        })

        if avg_score > best_score:
            best_score = avg_score
            best_param = {'alpha': alpha, 'solver': solver, 'tol': tol}

    best_params[label] = best_param
    best_scores[label] = best_score

    print(f"  🏆 Best: alpha={best_param['alpha']}, solver={best_param['solver']}, "
          f"tol={best_param['tol']}, R²={best_score:.6f}")

# Calculate overall OOF with best params
print("\n" + "="*70)
print("📈 Training with Best Parameters")
print("="*70)

oof_preds = [[] for _ in range(5)]
kf = KFold(n_splits=5, shuffle=True, random_state=42)

for i, label in enumerate(LABELS):
    print(f"\n📊 Target: {label}")
    targets_np = np.array(torch.cat(targets[i]))
    params = best_params[label]
    print(f"  Using: alpha={params['alpha']}, solver={params['solver']}, tol={params['tol']}")

    for fold, (train_idxs, val_idxs) in enumerate(kf.split(embeds_np)):
        X_train, X_val = embeds_np[train_idxs], embeds_np[val_idxs]
        y_train, y_val = targets_np[train_idxs], targets_np[val_idxs]

        reg = Ridge(**params, max_iter=10000)
        reg.fit(X_train, y_train)

        preds = reg.predict(X_val)
        oof_preds[i].extend(zip(val_idxs, preds))

# Calculate final OOF scores
oof_scores = {}
for i, label in enumerate(LABELS):
    sorted_preds = sorted(oof_preds[i], key=lambda x: x[0])
    pred_values = np.array([p for _, p in sorted_preds])
    true_values = np.array(torch.cat(targets[i]))

    ss_res = np.sum((true_values - pred_values) ** 2)
    ss_tot = np.sum((true_values - np.mean(true_values)) ** 2)
    r2 = 1 - (ss_res / ss_tot)
    oof_scores[label] = r2

weighted_r2 = sum(oof_scores[label] * WEIGHTS[label] for label in LABELS)

print("\n" + "="*70)
print("🎯 Final OOF Scores")
print("="*70)
for label in LABELS:
    print(f"  {label:20s}: R² = {oof_scores[label]:.6f}")
print(f"\n🎯 Overall OOF Score (Weighted R²): {weighted_r2:.6f}")
print("="*70)

# Save results
grid_results_df = pd.DataFrame(all_results)
grid_results_df.to_csv('grid_search_results.csv', index=False)

with open('exp103_results.txt', 'w') as f:
    f.write(f"Exp103: DINOv2 + Ridge Grid Search\n")
    f.write(f"=" * 70 + "\n\n")
    f.write(f"Best Parameters per Target:\n")
    for label in LABELS:
        params = best_params[label]
        f.write(f"  {label:20s}: alpha={params['alpha']}, solver={params['solver']}, "
                f"tol={params['tol']}, R²={oof_scores[label]:.6f}\n")
    f.write(f"\nOverall OOF Score: {weighted_r2:.6f}\n\n")
    f.write(f"Comparison:\n")
    f.write(f"  Exp101 (Ridge alpha=0.5): (baseline)\n")
    f.write(f"  Exp103 (grid search):     {weighted_r2:.6f}\n")

print("\n✅ Grid search results saved")

# Train final models
print("\n" + "="*70)
print("🎯 Training final models for submission...")
print("="*70)

final_regressors = [[None for _ in range(5)] for _ in range(5)]

for i, label in enumerate(LABELS):
    targets_np = np.array(torch.cat(targets[i]))
    params = best_params[label]

    for fold, (train_idxs, val_idxs) in enumerate(kf.split(embeds_np)):
        X_train = embeds_np[train_idxs]
        y_train = targets_np[train_idxs]

        reg = Ridge(**params, max_iter=10000)
        reg.fit(X_train, y_train)
        final_regressors[i][fold] = reg

# Test inference
print("\n📝 Running predictions on test data...")
test_df = pd.read_csv(os.path.join(ROOT, "test.csv"))
test_embeds = {}

for img_path in tqdm(test_df['image_path'].unique(), desc="  Extracting test features"):
    full_path = os.path.join(ROOT, img_path)
    with Image.open(full_path) as img:
        inputs = processor(images=img, return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        outputs = model(**inputs)
        img_name = os.path.splitext(os.path.basename(img_path))[0]
        test_embeds[img_name] = outputs.pooler_output.cpu()

predictions, sample_ids = [], []

for _, entry in test_df.iterrows():
    sample_id = entry['sample_id']
    img_name, target_name = sample_id.split('__')
    X_test = np.array(test_embeds[img_name])
    target_idx = target_mapping[target_name]

    fold_preds = [reg.predict(X_test) for reg in final_regressors[target_idx]]
    prediction = np.mean(fold_preds)
    predictions.append(max(0.0, prediction))
    sample_ids.append(sample_id)

submission = pd.DataFrame({'sample_id': sample_ids, 'target': predictions})
submission = submission.sort_values('sample_id').reset_index(drop=True)
submission.to_csv('submission.csv', index=False)
print(f"\n✅ Submission created: {len(submission)} rows")

print("\n" + "="*70)
print("✨ Exp103 Complete!")
print(f"🎯 FINAL OOF SCORE: {weighted_r2:.6f}")
print("="*70)

# Cleanup
del model
gc.collect()
if torch.cuda.is_available():
    torch.cuda.empty_cache()