#!/usr/bin/env python3
"""
Exp05: Improved Ensemble (70% DINOv2 + 20% ConvNext + 10% SigLIP)

Based on Exp03 results showing DINOv2 significantly outperforms other models:
- DINOv2 OOF: 0.6847
- SigLIP OOF: 0.5953
- ConvNext OOF: ~0.61 (estimated)

New weights: 70% DINOv2, 20% ConvNext, 10% SigLIP
"""

import pandas as pd
import numpy as np

print("="*70)
print("🚀 Exp05: Improved Ensemble (DINOv2-Heavy)")
print("="*70)
print("New weights: 70% DINOv2 + 20% ConvNext + 10% SigLIP")
print("="*70)

# Load OOF predictions
print("\n📦 Loading OOF predictions...")
dinov2_oof = pd.read_csv('/kaggle/input/dinov2-giant-simple-baseline/oof_dinov2_giant.csv')
siglip_oof = pd.read_csv('/kaggle/input/csiro-model-ensembles-lb-0-634/Model_2_SigLIP_pred.csv')
convnext_oof = pd.read_csv('/kaggle/input/csiro-model-ensembles-lb-0-634/Model_3_ConvNeXt_pred.csv')
print("✅ OOF predictions loaded")

# Rename columns for consistency
dinov2_oof = dinov2_oof.rename(columns={
    'Dry_Clover_g': 'Dry_Clover_g_dinov2',
    'Dry_Dead_g': 'Dry_Dead_g_dinov2',
    'Dry_Green_g': 'Dry_Green_g_dinov2',
    'Dry_Total_g': 'Dry_Total_g_dinov2',
    'GDM_g': 'GDM_g_dinov2'
})

siglip_oof = siglip_oof.rename(columns={
    'Dry_Clover_g': 'Dry_Clover_g_siglip',
    'Dry_Dead_g': 'Dry_Dead_g_siglip',
    'Dry_Green_g': 'Dry_Green_g_siglip',
    'Dry_Total_g': 'Dry_Total_g_siglip',
    'GDM_g': 'GDM_g_siglip'
})

convnext_oof = convnext_oof.rename(columns={
    'Dry_Clover_g': 'Dry_Clover_g_convnext',
    'Dry_Dead_g': 'Dry_Dead_g_convnext',
    'Dry_Green_g': 'Dry_Green_g_convnext',
    'Dry_Total_g': 'Dry_Total_g_convnext',
    'GDM_g': 'GDM_g_convnext'
})

# Merge all predictions
print("\n🔗 Merging predictions...")
merged = dinov2_oof.merge(siglip_oof, on='image_path')
merged = merged.merge(convnext_oof, on='image_path')
print(f"✅ Merged {len(merged)} images")

# Create improved ensemble (70-20-10)
print("\n📊 Creating improved ensemble...")
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}

ensemble_oof = merged[['image_path']].copy()
for label in LABELS:
    ensemble_oof[label] = (
        0.70 * merged[f'{label}_dinov2'] +
        0.20 * merged[f'{label}_convnext'] +
        0.10 * merged[f'{label}_siglip']
    )

# Save ensemble OOF
ensemble_oof.to_csv('oof_improved_ensemble.csv', index=False)
print("✅ Ensemble OOF saved")

# Calculate OOF score
print("\n" + "="*70)
print("📊 Calculating OOF Score...")
print("="*70)

# Load ground truth
ROOT = "/kaggle/input/csiro-biomass/"
train_df = pd.read_csv(ROOT + "train.csv")
train_wide = train_df.pivot_table(
    index='image_path',
    columns='target_name',
    values='target',
    aggfunc='first'
).reset_index()

# Merge with ground truth
merged_eval = pd.merge(ensemble_oof, train_wide, on='image_path', suffixes=('', '_true'))

# Extract predictions and ground truth
y_pred = merged_eval[LABELS].values
y_true = merged_eval[[f'{col}_true' for col in LABELS]].values

# Competition metric (weighted R²)
weights_array = np.array([WEIGHTS[l] for l in LABELS])
y_weighted_mean = np.average(y_true, weights=weights_array, axis=1).mean()
ss_res = np.average((y_true - y_pred)**2, weights=weights_array, axis=1).mean()
ss_tot = np.average((y_true - y_weighted_mean)**2, weights=weights_array, axis=1).mean()
oof_score = 1 - ss_res / ss_tot

print(f"\n🎯 OOF Score (Weighted R²): {oof_score:.6f}")
print("="*70)

# Save score
with open('exp05_oof_score.txt', 'w') as f:
    f.write(f"Exp05 Improved Ensemble OOF Score: {oof_score:.6f}\n")
    f.write(f"Weights: 70% DINOv2 + 20% ConvNext + 10% SigLIP\n")
    f.write(f"\nComparison:\n")
    f.write(f"  Original Ensemble (15-45-40): 0.6352\n")
    f.write(f"  Improved Ensemble (70-20-10): {oof_score:.6f}\n")
    f.write(f"  DINOv2 Alone: 0.6847\n")

print(f"\n📝 Comparison:")
print(f"  Original Ensemble (15-45-40): 0.6352")
print(f"  Improved Ensemble (70-20-10): {oof_score:.6f}")
print(f"  DINOv2 Alone: 0.6847")

# Create submission file
print("\n" + "="*70)
print("📝 Creating Submission File...")
print("="*70)

# Load test predictions
dinov2_test = pd.read_csv('/kaggle/input/dinov2-giant-simple-baseline/submission.csv')
siglip_test = pd.read_csv('/kaggle/input/csiro-model-ensembles-lb-0-634/Model_2_SigLIP.csv')
convnext_test = pd.read_csv('/kaggle/input/csiro-model-ensembles-lb-0-634/Model_3_ConvNeXt.csv')

# Create ensemble submission
submission = dinov2_test[['sample_id']].copy()
submission['target'] = (
    0.70 * dinov2_test['target'] +
    0.20 * convnext_test['target'] +
    0.10 * siglip_test['target']
)

# Ensure non-negative
submission['target'] = submission['target'].clip(lower=0.0)

# Save submission
submission.to_csv('submission.csv', index=False)
print(f"✅ Submission file created: {len(submission)} rows")

print(f"\n📋 First 10 rows:")
print(submission.head(10))

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