#!/usr/bin/env python3
"""
Exp106: High OOF Ensemble - DINOv2優先アンサンブル

戦略:
- 元のアンサンブル (LB 0.634): DINOv2(15%) + SigLIP(40%) + ConvNext(45%)
- 新アンサンブル: DINOv2(70%) + SigLIP(15%) + ConvNext(15%)
- 理由: DINOv2単体のOOF 0.6847が最強、重みを逆転させる

期待OOF: 0.68-0.70
期待Public LB: 0.67-0.69 (元の0.634を大きく上回る)
"""

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
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("🚀 Exp106: High OOF Ensemble (DINOv2-focused)")
print("="*70)
print(f"Device: {DEVICE}")
print("New weights: DINOv2(70%) + SigLIP(15%) + ConvNext(15%)")
print("Original weights: DINOv2(15%) + SigLIP(40%) + ConvNext(45%)")
print("Expected improvement: OOF 0.68-0.70, Public LB > 0.67")
print("="*70)

# Load existing predictions from Exp01 (3-model ensemble)
print("\n📂 Loading existing model predictions from Exp01...")
EXP01_PATH = "/kaggle/input/csiro-model-ensembles-lb-0-634"

# Model 1: DINOv2 predictions
model1_train = pd.read_csv(os.path.join(EXP01_PATH, "oof_model1_dinov2.csv"))
model1_test = pd.read_csv(os.path.join(EXP01_PATH, "submission_model1_dinov2.csv"))

# Model 2: ConvNext predictions
model2_train = pd.read_csv(os.path.join(EXP01_PATH, "oof_model2_convnext.csv"))
model2_test = pd.read_csv(os.path.join(EXP01_PATH, "submission_model2_convnext.csv"))

# Model 3: SigLIP predictions
model3_train = pd.read_csv(os.path.join(EXP01_PATH, "oof_model3_siglip.csv"))
model3_test = pd.read_csv(os.path.join(EXP01_PATH, "submission_model3_siglip.csv"))

print(f"✅ Loaded predictions:")
print(f"  Model 1 (DINOv2): {len(model1_train)} train samples")
print(f"  Model 2 (ConvNext): {len(model2_train)} train samples")
print(f"  Model 3 (SigLIP): {len(model3_train)} train samples")

# Load training labels for OOF calculation
train_df = pd.read_csv(os.path.join(ROOT, "train.csv"))

# Create ensemble with NEW weights
print("\n🔄 Creating ensemble predictions...")
print("New ensemble weights:")
print("  - DINOv2:   70% (was 15%)")
print("  - SigLIP:   15% (was 40%)")
print("  - ConvNext: 15% (was 45%)")

# NEW WEIGHTS (DINOv2-focused)
W1 = 0.70  # DINOv2 (最強モデル)
W2 = 0.15  # ConvNext
W3 = 0.15  # SigLIP

# Training predictions (OOF)
oof_ensemble = model1_train.copy()
oof_ensemble['target'] = (
    W1 * model1_train['target'] +
    W2 * model2_train['target'] +
    W3 * model3_train['target']
)

# Test predictions
test_ensemble = model1_test.copy()
test_ensemble['target'] = (
    W1 * model1_test['target'] +
    W2 * model2_test['target'] +
    W3 * model3_test['target']
)

# Calculate OOF scores per target
print("\n📊 Calculating OOF scores per target...")
oof_scores = {}

for label in LABELS:
    # Get samples for this label
    label_samples = train_df[train_df['sample_id'].str.contains(f'__{label}')]

    # Merge with ensemble predictions
    oof_label = oof_ensemble[oof_ensemble['sample_id'].isin(label_samples['sample_id'])]

    # Calculate R²
    true_values = label_samples.set_index('sample_id').loc[oof_label['sample_id'], 'target'].values
    pred_values = oof_label['target'].values

    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
    print(f"  {label}: R² = {r2:.6f}")

# Calculate weighted OOF
weighted_oof = sum(oof_scores[label] * WEIGHTS[label] for label in LABELS)

print(f"\n" + "="*70)
print(f"🎯 Overall Weighted OOF Score: {weighted_oof:.6f}")
print(f"="*70)

# Compare with original ensemble
print(f"\n📈 Comparison with Original Ensemble:")
print(f"  Original (15/40/45): Public LB 0.634 (OOF unknown)")
print(f"  New (70/15/15):      OOF {weighted_oof:.6f}")
print(f"  Expected Public LB:  0.67-0.69")

# Save OOF predictions
oof_ensemble.to_csv('oof_exp106_high_oof_ensemble.csv', index=False)
print(f"\n✅ Saved OOF predictions: oof_exp106_high_oof_ensemble.csv")

# Save test predictions (submission)
test_ensemble.to_csv('submission.csv', index=False)
print(f"✅ Saved submission: submission.csv")

# Save results summary
with open('exp106_results.txt', 'w') as f:
    f.write(f"Exp106: High OOF Ensemble (DINOv2-focused)\n")
    f.write(f"=" * 70 + "\n\n")
    f.write(f"Ensemble Weights:\n")
    f.write(f"  DINOv2:   {W1:.0%} (was 15%)\n")
    f.write(f"  ConvNext: {W2:.0%} (was 45%)\n")
    f.write(f"  SigLIP:   {W3:.0%} (was 40%)\n\n")
    f.write(f"OOF Scores per Target:\n")
    for label in LABELS:
        f.write(f"  {label}: {oof_scores[label]:.6f}\n")
    f.write(f"\nOverall Weighted OOF: {weighted_oof:.6f}\n\n")
    f.write(f"Comparison:\n")
    f.write(f"  Original (15/40/45): Public LB 0.634\n")
    f.write(f"  New (70/15/15):      OOF {weighted_oof:.6f}\n")
    f.write(f"  Expected Public LB:  0.67-0.69\n")

print(f"✅ Saved results: exp106_results.txt")
print(f"\n🎯 Exp106 Complete!")
print(f"Expected improvement: +{(weighted_oof - 0.634) / 0.634 * 100:.1f}% from original ensemble")

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