#!/usr/bin/env python3
"""
Exp14: SigLIP Single Model Evaluation
Extract features from SigLIP and train CatBoost/GradientBoosting ensemble
Calculate OOF score and generate submission
"""

import warnings
warnings.filterwarnings('ignore')

from pathlib import Path
from tqdm.auto import tqdm
import json
from copy import deepcopy
import polars as pl
import numpy as np
import pandas as pd
import os

import torch
from PIL import Image
from transformers import AutoProcessor, AutoImageProcessor, AutoModel

from sklearn.model_selection import GroupKFold
from sklearn.ensemble import GradientBoostingRegressor
import catboost

# ============================================================================
# Configuration
# ============================================================================

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"Using device: {device}")

data_path = Path('/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,
}

# ============================================================================
# Competition Metric
# ============================================================================

def competition_metric(y_true, y_pred) -> float:
    """Calculate competition's official weighted R2 score"""
    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()

    return 1 - ss_res / ss_tot

# ============================================================================
# 1. Load SigLIP Model and Extract Features
# ============================================================================

print("="*70)
print("Exp14: SigLIP Single Model")
print("="*70)

print("\nLoading SigLIP model...")
model_name = "/kaggle/input/google-siglip-so400m-patch14-384/transformers/default/1/"
model = AutoModel.from_pretrained(model_name)
model = model.to(device)
model.eval()

processor = AutoImageProcessor.from_pretrained(model_name)
print("Model loading complete.")

# ============================================================================
# 2. Load and Preprocess Data
# ============================================================================

print("\nProcessing training data...")
train = pl.read_csv(data_path / 'train.csv')
df = (
    train
    .with_columns([
        pl.when(pl.col('target_name') == label).then(pl.col('target')).alias(label)
        for label in labels
    ])
    .group_by('image_path')
    .agg([
        pl.col(label).mean()
        for label in labels
    ] + [
        pl.concat_str(["Sampling_Date", "State"], separator=" ").alias("group").first()
    ])
    .sort('image_path')
)

print("Processing test data...")
test = pl.read_csv(data_path / 'test.csv')
df_test = (
    test
    .group_by('image_path')
    .len()
    .sort('image_path')
)

# ============================================================================
# 3. Extract Features
# ============================================================================

def compute_features(images: list, save_path: str):
    """Extract features from images using SigLIP model"""
    batch_size = 20
    with torch.no_grad(), open(save_path, 'w') as f:
        for i in tqdm(range(0, len(images), batch_size), desc=f"Extracting {save_path}"):
            batch_paths = images[i:i + batch_size]
            batch = [Image.open(data_path / p) for p in batch_paths]
            inputs = processor(images=batch, return_tensors="pt").to(model.device)
            features = model.get_image_features(**inputs)
            for line in features:
                data = {f'x_{j}': line[j].item() for j in range(len(line))}
                f.write(json.dumps(data) + '\n')

compute_features(df['image_path'], 'features_train.ndjson')
compute_features(df_test['image_path'], 'features_test.ndjson')
print("Feature extraction complete.")

# ============================================================================
# 4. Combine Features with Data
# ============================================================================

responses_train = pl.read_ndjson('features_train.ndjson')
responses_test = pl.read_ndjson('features_test.ndjson')
df_aug = pl.concat([df, responses_train], how='horizontal')
df_test_aug = pl.concat([df_test, responses_test], how='horizontal')

feature_columns = sorted(responses_train.columns)

# ============================================================================
# 5. Cross-Validation with OOF Predictions
# ============================================================================

def cross_validate_with_oof(model, data, data_test, x_columns, random_state=42):
    """Perform GroupKFold CV and return OOF predictions"""
    X = data.select(x_columns).to_numpy()
    X_test = data_test.select(x_columns).to_numpy()
    y_true = data.select(labels).to_numpy()

    y_pred_oof = np.zeros_like(y_true)
    y_pred_test = np.zeros([len(X_test), len(labels)])

    n_splits = 5
    kf = GroupKFold(n_splits=n_splits)
    groups = data.select('group')

    print("\nStarting 5-Fold Cross-Validation...")
    for i, (train_index, val_index) in enumerate(kf.split(X, groups=groups)):
        for l in range(len(labels)):
            m = deepcopy(model)
            m.fit(X[train_index], y_true[train_index, l])
            y_pred_oof[val_index, l] = m.predict(X[val_index]).clip(0)
            y_pred_test[:, l] += m.predict(X_test).clip(0) / n_splits

        score = competition_metric(y_true[val_index], y_pred_oof[val_index])
        print(f'  Fold {i}: Score = {score:.6f}')

    full_cv_score = competition_metric(y_true, y_pred_oof)
    print(f'\n🎯 Full OOF Score: {full_cv_score:.6f}')

    return y_pred_oof, y_pred_test, full_cv_score

# ============================================================================
# 6. Train Ensemble Models
# ============================================================================

print("\n" + "="*70)
print("Training GradientBoostingRegressor...")
print("="*70)
oof_pred_gb, pred_test_gb, score_gb = cross_validate_with_oof(
    GradientBoostingRegressor(random_state=42),
    df_aug,
    df_test_aug,
    feature_columns
)

print("\n" + "="*70)
print("Training CatBoostRegressor...")
print("="*70)
oof_pred_cb, pred_test_cb, score_cb = cross_validate_with_oof(
    catboost.CatBoostRegressor(verbose=False, iterations=100, random_state=42),
    df_aug,
    df_test_aug,
    feature_columns
)

# ============================================================================
# 7. Ensemble and Save OOF Predictions
# ============================================================================

print("\n" + "="*70)
print("Creating Ensemble (GB + CatBoost)")
print("="*70)

oof_pred_ensemble = (oof_pred_gb + oof_pred_cb) / 2
pred_test_ensemble = (pred_test_gb + pred_test_cb) / 2

# Calculate ensemble OOF score
y_true = df_aug.select(labels).to_numpy()
ensemble_score = competition_metric(y_true, oof_pred_ensemble)
print(f"\n🎯 Ensemble OOF Score: {ensemble_score:.6f}")

# Save OOF predictions
oof_df = df_aug.select(['image_path']).to_pandas()
oof_df[labels] = oof_pred_ensemble
oof_df.to_csv('oof_exp14_siglip.csv', index=False)
print("\n✅ OOF predictions saved: oof_exp14_siglip.csv")

# ============================================================================
# 8. Create Submission File
# ============================================================================

pred_with_id = pl.concat([
    df_test.select("image_path"),
    pl.DataFrame(pred_test_ensemble, schema=labels),
], how='horizontal')

pred_save = (
    test
    .join(pred_with_id, on='image_path')
    .with_columns(
        pl.coalesce(*[
            pl.when(pl.col('target_name') == col).then(pl.col(col))
            for col in labels
        ]).alias('target')
    )
    .select('sample_id', 'target')
)

pred_save.write_csv('submission.csv')
print("✅ Submission file saved: submission.csv")

# ============================================================================
# 9. Summary Report
# ============================================================================

print("\n" + "="*70)
print("📊 EXP14 SUMMARY")
print("="*70)
print(f"Model: SigLIP (google/siglip-so400m-patch14-384)")
print(f"Regression: GradientBoosting + CatBoost Ensemble")
print(f"")
print(f"GradientBoosting OOF: {score_gb:.6f}")
print(f"CatBoost OOF:         {score_cb:.6f}")
print(f"Ensemble OOF:         {ensemble_score:.6f}")
print(f"")
print(f"Expected Public LB:   ~0.595 (based on Exp03)")
print("="*70)

# Save summary
summary = {
    'experiment': 'exp14_siglip_single',
    'model': 'SigLIP + GradientBoosting/CatBoost Ensemble',
    'oof_score': ensemble_score,
    'gb_score': score_gb,
    'catboost_score': score_cb,
    'num_folds': 5,
}

import json
with open('exp14_summary.json', 'w') as f:
    json.dump(summary, f, indent=2)

with open('exp14_oof_score.txt', 'w') as f:
    f.write(f"{ensemble_score:.6f}\n")

print("\n✨ Exp14 completed successfully!")