#!/usr/bin/env python3
"""
Exp104: DINOv2 + Lasso & Ridge Ensemble

Exp101-103の結果を待たずに、Lasso+Ridgeの両方を使ったアンサンブル実験。

アプローチ:
1. Lassoモデル（alpha=0.5）でOOF予測
2. Ridgeモデル（alpha=0.5）でOOF予測
3. 両方の予測を平均してアンサンブル

仮説: LassoとRidgeは異なる特徴の使い方をするため、
      アンサンブルすることで相補的な予測ができる。

期待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 Lasso, Ridge
from sklearn.model_selection import KFold
from PIL import Image
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}

print("="*70)
print("🚀 Exp104: DINOv2 + Lasso & Ridge Ensemble")
print("="*70)
print(f"Device: {DEVICE}")
print("Baseline (Exp08 Lasso): alpha=0.5, OOF 0.6868")
print("This experiment: Ensemble of Lasso + Ridge")
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))

# 5-Fold Cross-Validation with both Lasso and Ridge
print("\n" + "="*70)
print("🎯 Training Lasso & Ridge Ensemble")
print("="*70)

oof_preds_lasso = [[] for _ in range(5)]
oof_preds_ridge = [[] 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]))

    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 Lasso
        lasso = Lasso(alpha=0.5, max_iter=10000)
        lasso.fit(X_train, y_train)
        preds_lasso = lasso.predict(X_val)
        oof_preds_lasso[i].extend(zip(val_idxs, preds_lasso))

        # Train Ridge
        ridge = Ridge(alpha=0.5, max_iter=10000)
        ridge.fit(X_train, y_train)
        preds_ridge = ridge.predict(X_val)
        oof_preds_ridge[i].extend(zip(val_idxs, preds_ridge))

# Calculate OOF scores for each model
print("\n" + "="*70)
print("📈 Out-of-Fold Scores (Individual Models)")
print("="*70)

# Lasso scores
print("\n🔷 Lasso:")
oof_scores_lasso = {}
for i, label in enumerate(LABELS):
    sorted_preds = sorted(oof_preds_lasso[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_lasso[label] = r2
    print(f"  {label:20s}: R² = {r2:.6f}")

weighted_r2_lasso = sum(oof_scores_lasso[label] * WEIGHTS[label] for label in LABELS)
print(f"\n  Overall Lasso OOF: {weighted_r2_lasso:.6f}")

# Ridge scores
print("\n🔶 Ridge:")
oof_scores_ridge = {}
for i, label in enumerate(LABELS):
    sorted_preds = sorted(oof_preds_ridge[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_ridge[label] = r2
    print(f"  {label:20s}: R² = {r2:.6f}")

weighted_r2_ridge = sum(oof_scores_ridge[label] * WEIGHTS[label] for label in LABELS)
print(f"\n  Overall Ridge OOF: {weighted_r2_ridge:.6f}")

# Ensemble: Average of Lasso and Ridge
print("\n" + "="*70)
print("🎯 Ensemble (Lasso + Ridge Average)")
print("="*70)

oof_scores_ensemble = {}
for i, label in enumerate(LABELS):
    sorted_lasso = sorted(oof_preds_lasso[i], key=lambda x: x[0])
    sorted_ridge = sorted(oof_preds_ridge[i], key=lambda x: x[0])

    pred_lasso = np.array([p for _, p in sorted_lasso])
    pred_ridge = np.array([p for _, p in sorted_ridge])
    pred_ensemble = (pred_lasso + pred_ridge) / 2

    true_values = np.array(torch.cat(targets[i]))

    ss_res = np.sum((true_values - pred_ensemble) ** 2)
    ss_tot = np.sum((true_values - np.mean(true_values)) ** 2)
    r2 = 1 - (ss_res / ss_tot)
    oof_scores_ensemble[label] = r2
    print(f"  {label:20s}: R² = {r2:.6f}")

weighted_r2_ensemble = sum(oof_scores_ensemble[label] * WEIGHTS[label] for label in LABELS)
print(f"\n🎯 Overall Ensemble OOF Score: {weighted_r2_ensemble:.6f}")
print("="*70)

# Save results
with open('exp104_results.txt', 'w') as f:
    f.write(f"Exp104: DINOv2 + Lasso & Ridge Ensemble\n")
    f.write(f"=" * 70 + "\n\n")
    f.write(f"Individual Model Scores:\n")
    f.write(f"  Lasso OOF:    {weighted_r2_lasso:.6f}\n")
    f.write(f"  Ridge OOF:    {weighted_r2_ridge:.6f}\n\n")
    f.write(f"Ensemble Score:\n")
    f.write(f"  Ensemble OOF: {weighted_r2_ensemble:.6f}\n\n")
    f.write(f"Comparison:\n")
    f.write(f"  Exp08 (Lasso only):       0.6868\n")
    f.write(f"  Exp104 (Lasso+Ridge):     {weighted_r2_ensemble:.6f}\n")
    f.write(f"  Improvement:              {(weighted_r2_ensemble - 0.6868):.6f} ({(weighted_r2_ensemble - 0.6868)/0.6868*100:+.2f}%)\n")

print("\n✅ Results saved")

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

final_lasso = [[None for _ in range(5)] for _ in range(5)]
final_ridge = [[None for _ in range(5)] for _ in range(5)]

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

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

        lasso = Lasso(alpha=0.5, max_iter=10000)
        lasso.fit(X_train, y_train)
        final_lasso[i][fold] = lasso

        ridge = Ridge(alpha=0.5, max_iter=10000)
        ridge.fit(X_train, y_train)
        final_ridge[i][fold] = ridge

# 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]

    # Ensemble prediction
    lasso_preds = [reg.predict(X_test) for reg in final_lasso[target_idx]]
    ridge_preds = [reg.predict(X_test) for reg in final_ridge[target_idx]]

    pred_lasso = np.mean(lasso_preds)
    pred_ridge = np.mean(ridge_preds)
    prediction = (pred_lasso + pred_ridge) / 2

    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("✨ Exp104 Complete!")
print(f"🎯 FINAL OOF SCORE: {weighted_r2_ensemble:.6f}")
print(f"   Lasso: {weighted_r2_lasso:.6f}")
print(f"   Ridge: {weighted_r2_ridge:.6f}")
print("="*70)

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