# %% [code]
# ============================================================================
# FIX: Prevent protobuf conflicts BEFORE any imports
# ============================================================================

import sys, os, warnings, logging

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'
os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] = 'true'

for mod in list(sys.modules.keys()):
    if 'protobuf' in mod or 'google' in mod:
        del sys.modules[mod]

warnings.filterwarnings('ignore')
logging.getLogger().setLevel(logging.ERROR)
print("✓ Environment configured")

# %% [code]
# ============================================================================
# SETUP: Path Configuration
# ============================================================================

from pathlib import Path

print("\n" + "="*70)
print("KAGGLE ENVIRONMENT SETUP")
print("="*70 + "\n")

INPUT_DIR = Path('/kaggle/input')
WORKING_DIR = Path('/kaggle/working')

datasets = [d for d in INPUT_DIR.iterdir() if d.is_dir()]
print("Available datasets:")
for dataset in datasets:
    print(f"  - {dataset.name}")

BASE_PATH = None
for d in datasets:
    if any(x in d.name.lower() for x in ['biomass', 'csiro', 'pasture', 'image2biomass']):
        BASE_PATH = d
        break

if BASE_PATH is None:
    BASE_PATH = [d for d in datasets if 'dino' not in d.name.lower()][0]

print(f"\n✓ Found competition: {BASE_PATH.name}")

TRAIN_CSV = BASE_PATH / 'train.csv'
TEST_CSV = BASE_PATH / 'test.csv'
SUBMISSION_TEMPLATE = BASE_PATH / 'sample_submission.csv'

MODELS_DIR = WORKING_DIR / 'models'
MODELS_DIR.mkdir(exist_ok=True, parents=True)

print(f"\nVerifying files:")
print(f"  ✓ train.csv: {TRAIN_CSV.exists()}")
print(f"  ✓ test.csv: {TEST_CSV.exists()}")
print(f"  ✓ sample_submission.csv: {SUBMISSION_TEMPLATE.exists()}")

print("\n" + "="*70 + "\n")

# %% [code]
# ============================================================================
# IMPORTS
# ============================================================================

import pandas as pd
import numpy as np
from PIL import Image
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import r2_score, mean_squared_error

np.random.seed(42)
torch.manual_seed(42)
print("✓ Imports successful")

# %% [code]
# ============================================================================
# LOAD DATA
# ============================================================================

print("\n" + "="*70)
print("LOADING DATA")
print("="*70 + "\n")

train_df = pd.read_csv(TRAIN_CSV)
test_df = pd.read_csv(TEST_CSV)

print(f"Train samples: {len(train_df)}")
print(f"Test samples: {len(test_df)}")
print(f"Targets: {train_df['target_name'].unique().tolist()}")

# %% [code]
# ============================================================================
# DEVICE SETUP
# ============================================================================

print("\n" + "="*70)
print("DEVICE SETUP")
print("="*70 + "\n")

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

# %% [code]
# ============================================================================
# LOAD DINO V2 MODEL
# ============================================================================

print("="*70)
print("LOADING DINO V2 MODEL")
print("="*70 + "\n")

from transformers import AutoImageProcessor, AutoModel

dino_candidates = list(Path('/kaggle/input').rglob('preprocessor_config.json'))
if not dino_candidates:
    raise RuntimeError("DINOv2 model not found. Please add DINOv2 to Kaggle inputs.")
dino_path = str(dino_candidates[0].parent)

image_processor = AutoImageProcessor.from_pretrained(dino_path, local_files_only=True)
dino_model = AutoModel.from_pretrained(dino_path, local_files_only=True)
dino_model = dino_model.to(device)
dino_model.eval()
print(f"✓ Loaded DINOv2 from {dino_path}")

# %% [code]
# ============================================================================
# FEATURE EXTRACTION FUNCTION (Improved)
# ============================================================================

def extract_dino_features(image_path):
    """
    Extracts attention-weighted DINOv2 features.
    """
    try:
        img = Image.open(image_path).convert('RGB')
        inputs = image_processor(img, return_tensors="pt").to(device)
        with torch.no_grad():
            outputs = dino_model(**inputs)
        features = outputs.last_hidden_state.mean(dim=1)
        features = F.normalize(features, dim=1)
        return features.cpu().numpy().flatten()
    except Exception as e:
        return None

print("✓ Feature extraction function ready")

# %% [code]
# ============================================================================
# EXTRACT FEATURES FOR TRAINING
# ============================================================================

print("\n" + "="*70)
print("EXTRACTING TRAIN FEATURES")
print("="*70 + "\n")

train_features = {}
for img_path in tqdm(train_df['image_path'].unique(), desc="Train images"):
    full_path = BASE_PATH / img_path
    if full_path.exists():
        feat = extract_dino_features(str(full_path))
        if feat is not None:
            train_features[img_path] = feat

print(f"✓ Extracted {len(train_features)} image embeddings")

# %% [code]
# ============================================================================
# EXTRACT FEATURES FOR TEST
# ============================================================================

print("\n" + "="*70)
print("EXTRACTING TEST FEATURES")
print("="*70 + "\n")

test_features = {}
for img_path in tqdm(test_df['image_path'].unique(), desc="Test images"):
    full_path = BASE_PATH / img_path
    if full_path.exists():
        feat = extract_dino_features(str(full_path))
        if feat is not None:
            test_features[img_path] = feat

print(f"✓ Extracted {len(test_features)} test embeddings")

# %% [code]
# ============================================================================
# PREPARE TRAINING DATA
# ============================================================================

X, y, targets = [], [], []

for _, row in train_df.iterrows():
    if row['image_path'] in train_features:
        X.append(train_features[row['image_path']])
        y.append(row['target'])
        targets.append(row['target_name'])

X = np.array(X)
y = np.array(y)
target_names = np.array(targets)

scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)
print(f"✓ X_scaled shape: {X_scaled.shape}")

# Encode targets
unique_targets = np.unique(target_names)
target_to_idx = {t: i for i, t in enumerate(unique_targets)}
y_targets = np.array([target_to_idx[t] for t in target_names])

print(f"✓ Encoded {len(unique_targets)} targets: {unique_targets}")

# %% [code]
# ============================================================================
# DEFINE MLP REGRESSOR
# ============================================================================

class BiomassMLP(nn.Module):
    def __init__(self, input_dim=768, hidden_dim=256, output_dim=5):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim, 128),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(128, output_dim)
        )

    def forward(self, x):
        return self.mlp(x)

target_weights = torch.tensor([0.1, 0.1, 0.1, 0.2, 0.5], device=device)

def weighted_mse_loss(pred, true):
    mse = (pred - true) ** 2
    return (mse * target_weights).mean()

print("✓ BiomassMLP and weighted loss defined")

# %% [code]
# ============================================================================
# TRAINING
# ============================================================================

print("\n" + "="*70)
print("TRAINING MODEL (Algorithmic improvements only)")
print("="*70 + "\n")

# Prepare data in PyTorch tensors
X_tensor = torch.tensor(X_scaled, dtype=torch.float32).to(device)
y_tensor = torch.zeros((len(y), 5), dtype=torch.float32).to(device)

# Map scalar y to corresponding column in y_tensor
for i, (val, tname) in enumerate(zip(y, target_names)):
    y_tensor[i, target_to_idx[tname]] = val

# Split into train/validation sets
train_idx, val_idx = train_test_split(np.arange(len(y)), test_size=0.2, random_state=42)
X_train, X_val = X_tensor[train_idx], X_tensor[val_idx]
y_train, y_val = y_tensor[train_idx], y_tensor[val_idx]

# Initialize model and training utilities
model = BiomassMLP(input_dim=X_tensor.shape[1], output_dim=5).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
grad_scaler = torch.cuda.amp.GradScaler()  # renamed to avoid conflict with StandardScaler

num_epochs = 15
batch_size = 128

for epoch in range(num_epochs):
    model.train()
    total_loss = 0.0

    for i in range(0, len(X_train), batch_size):
        xb = X_train[i:i+batch_size]
        yb = y_train[i:i+batch_size]
        optimizer.zero_grad()

        with torch.cuda.amp.autocast():
            preds = model(xb)
            loss = weighted_mse_loss(preds, yb)

        grad_scaler.scale(loss).backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        grad_scaler.step(optimizer)
        grad_scaler.update()
        scheduler.step()

        total_loss += loss.item()

    # Validation step
    model.eval()
    with torch.no_grad():
        val_preds = model(X_val)
        val_loss = weighted_mse_loss(val_preds, y_val).item()

    print(f"Epoch {epoch+1}/{num_epochs} | Train Loss: {total_loss:.4f} | Val Loss: {val_loss:.4f}")

print("✓ Training complete")

# %% [code]
# ============================================================================
# INFERENCE
# ============================================================================

print("\n" + "="*70)
print("INFERENCE AND SUBMISSION")
print("="*70 + "\n")

# Prepare test features in the same order as test_df
predictions = []

for _, row in tqdm(test_df.iterrows(), total=len(test_df), desc="Predicting"):
    img_path = row['image_path']
    target_name = row['target_name']

    if img_path in test_features:
        # Use the StandardScaler (kept as 'scaler') for feature normalization
        feat = scaler.transform(test_features[img_path].reshape(1, -1))
        feat_tensor = torch.tensor(feat, dtype=torch.float32).to(device)

        with torch.no_grad(), torch.cuda.amp.autocast():
            preds = model(feat_tensor).cpu().numpy().flatten()

        pred_val = preds[target_to_idx[target_name]]
        predictions.append({'sample_id': row['sample_id'], 'target': float(pred_val)})
    else:
        predictions.append({'sample_id': row['sample_id'], 'target': 0.0})

# Save predictions
preds_df = pd.DataFrame(predictions)
submission_path = WORKING_DIR / 'submission.csv'
preds_df.to_csv(submission_path, index=False)

print(f"✓ Submission saved: {submission_path}")
print(f"✓ Total predictions: {len(preds_df)}")
print(f"✓ Mean: {preds_df['target'].mean():.4f}, Std: {preds_df['target'].std():.4f}")


# %% [code]
# ============================================================================
# CLEANUP (Kaggle compliance)
# ============================================================================

import shutil
if os.path.exists("models"):
    shutil.rmtree("models")
print("✓ Models folder deleted. Notebook ready for submission.")