#!/usr/bin/env python3
"""
Exp105: DINOv2 + ElasticNet (L1 + L2 Regularization)

ElasticNetは、LassoとRidgeの良いとこ取り:
- L1正則化（Lasso）: 特徴選択
- L2正則化（Ridge）: 係数の安定化

正則化項: α * (ρ * L1 + (1-ρ) * L2)

where:
- α (alpha): 正則化の強さ
- ρ (l1_ratio): L1とL2のバランス
  - ρ=1.0: Lasso
  - ρ=0.0: Ridge
  - ρ=0.5: 50/50のバランス

期待OOF: 0.69-0.71
"""

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 ElasticNet
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],
    'l1_ratio': [0.3, 0.5, 0.7]  # L1とL2のバランス
}

print("="*70)
print("🚀 Exp105: DINOv2 + ElasticNet Grid Search")
print("="*70)
print(f"Device: {DEVICE}")
print("ElasticNet = Lasso + Ridge の良いとこ取り")
print("\nGrid Search Parameters:")
print(f"  alpha:     {PARAM_GRID['alpha']}")
print(f"  l1_ratio:  {PARAM_GRID['l1_ratio']}")
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 ElasticNet 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, l1_ratio = 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 ElasticNet
            reg = ElasticNet(alpha=alpha, l1_ratio=l1_ratio, 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,
            'l1_ratio': l1_ratio,
            'r2': avg_score
        })

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

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

    l1_pct = int(best_param['l1_ratio'] * 100)
    l2_pct = 100 - l1_pct
    print(f"  🏆 Best: alpha={best_param['alpha']}, "
          f"L1={l1_pct}%/L2={l2_pct}%, R²={best_score:.6f}")

# Calculate overall OOF with best params
print("\n" + "="*70)
print("📈 Training with Best ElasticNet 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]
    l1_pct = int(params['l1_ratio'] * 100)
    l2_pct = 100 - l1_pct
    print(f"  Using: alpha={params['alpha']}, L1={l1_pct}%/L2={l2_pct}%")

    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 = ElasticNet(**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('exp105_results.txt', 'w') as f:
    f.write(f"Exp105: DINOv2 + ElasticNet 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]
        l1_pct = int(params['l1_ratio'] * 100)
        l2_pct = 100 - l1_pct
        f.write(f"  {label:20s}: alpha={params['alpha']}, "
                f"L1={l1_pct}%/L2={l2_pct}%, 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"  Exp08 (Lasso):       0.6868\n")
    f.write(f"  Exp105 (ElasticNet): {weighted_r2:.6f}\n")
    f.write(f"  Improvement:         {(weighted_r2 - 0.6868):.6f} ({(weighted_r2 - 0.6868)/0.6868*100:+.2f}%)\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 = ElasticNet(**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("✨ Exp105 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()