"""
Weighted Log-Space Model for Biomass Prediction
================================================
Properly aligned with competition evaluation:
- Weighted loss matching competition weights
- Log transformation log(1+y)
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import timm
import pandas as pd
import numpy as np
import cv2
from pathlib import Path
from tqdm import tqdm

# Settings
DATA_DIR = Path("/kaggle/input/csiro-biomass")
WEIGHTS_DIR = Path("/kaggle/input/weighted-log-biomass-weights")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
IMAGE_SIZE = 384

print(f"Device: {DEVICE}")

# Simple 6-channel vegetation indices
def compute_veg_indices(img):
    img_float = img.astype(np.float32) / 255.0
    R, G, B = img_float[:,:,0], img_float[:,:,1], img_float[:,:,2]
    eps = 1e-6
    exg = np.clip((2*G - R - B + 1) / 2, 0, 1)
    vari = np.clip(((G - R) / (G + R - B + eps) + 1) / 2, 0, 1)
    gli = np.clip(((2*G - R - B) / (2*G + R + B + eps) + 1) / 2, 0, 1)
    return np.stack([R, G, B, exg, vari, gli], axis=-1)

MEAN = np.array([0.485, 0.456, 0.406, 0.5, 0.5, 0.5])
STD = np.array([0.229, 0.224, 0.225, 0.25, 0.25, 0.25])

def process_image(img_path, tta_idx=0):
    """Process image with optional TTA"""
    img = cv2.imread(str(img_path))
    if img is None:
        raise ValueError(f"Could not read image: {img_path}")
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (IMAGE_SIZE, IMAGE_SIZE))

    # TTA augmentations
    if tta_idx == 1:
        img = np.fliplr(img).copy()
    elif tta_idx == 2:
        img = np.flipud(img).copy()
    elif tta_idx == 3:
        img = np.fliplr(np.flipud(img)).copy()
    elif tta_idx == 4:
        img = np.rot90(img, 1).copy()
    elif tta_idx == 5:
        img = np.rot90(img, 2).copy()
    elif tta_idx == 6:
        img = np.rot90(img, 3).copy()
    elif tta_idx == 7:
        img = np.fliplr(np.rot90(img, 1)).copy()

    img_6ch = compute_veg_indices(img)
    img_6ch = (img_6ch - MEAN) / STD
    img_6ch = np.transpose(img_6ch, (2, 0, 1)).astype(np.float32)
    return torch.tensor(img_6ch)

# Model architecture - same as training
class WeightedLogModel(nn.Module):
    def __init__(self, in_channels=6):
        super().__init__()
        self.backbone = timm.create_model("convnext_base.fb_in22k_ft_in1k_384", pretrained=False, num_classes=0, in_chans=3)

        # Modify first conv for 6 channels
        old_conv = self.backbone.stem[0]
        new_conv = nn.Conv2d(in_channels, old_conv.out_channels, old_conv.kernel_size, old_conv.stride, old_conv.padding)
        self.backbone.stem[0] = new_conv

        # Get feature dim
        with torch.no_grad():
            dummy = torch.zeros(1, in_channels, IMAGE_SIZE, IMAGE_SIZE)
            feat_dim = self.backbone(dummy).shape[1]

        # Predict 3 primary targets
        self.head = nn.Sequential(
            nn.Dropout(0.5),
            nn.Linear(feat_dim, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 3)
        )

    def forward(self, x):
        feat = self.backbone(x)
        out = self.head(feat)
        primary = F.softplus(out)  # [Green, Dead, Clover]

        # Derive other targets
        green, dead, clover = primary[:, 0], primary[:, 1], primary[:, 2]
        gdm = green + clover
        total = green + dead + clover

        return torch.stack([green, dead, clover, gdm, total], dim=1)

# Load model
print("Loading model...")
model = WeightedLogModel().to(DEVICE)
ckpt = torch.load(WEIGHTS_DIR / "weighted_log_final.pth", map_location=DEVICE, weights_only=False)
model.load_state_dict(ckpt["model_state_dict"])
model.eval()
print(f"Model loaded: train_loss={ckpt['train_loss']:.4f}, CV mean={ckpt['cv_mean']:.4f}")

# Load test data
test_df = pd.read_csv(DATA_DIR / "test.csv")
print(f"\nTest data shape: {test_df.shape}")
print(f"Columns: {test_df.columns.tolist()}")

# Get unique images
unique_images = test_df["image_path"].unique()
print(f"\nProcessing {len(unique_images)} test images with 8x TTA...")

NUM_TTA = 8

# Run inference
image_predictions = {}
with torch.no_grad():
    for img_path in tqdm(unique_images, desc="Inference"):
        full_path = DATA_DIR / img_path

        # Collect predictions from TTA
        tta_preds = []
        for tta_idx in range(NUM_TTA):
            img_tensor = process_image(full_path, tta_idx).unsqueeze(0).to(DEVICE)
            pred = model(img_tensor)
            tta_preds.append(pred.cpu().numpy()[0])

        # Average TTA predictions
        avg_pred = np.mean(tta_preds, axis=0)

        # Store all 5 predictions
        image_predictions[img_path] = {
            "Dry_Green_g": float(avg_pred[0]),
            "Dry_Dead_g": float(avg_pred[1]),
            "Dry_Clover_g": float(avg_pred[2]),
            "GDM_g": float(avg_pred[3]),
            "Dry_Total_g": float(avg_pred[4]),
        }

# Create submission
submission = []
for _, row in test_df.iterrows():
    sample_id = row["sample_id"]
    img_path = row["image_path"]
    target_name = row["target_name"]

    pred_value = image_predictions[img_path][target_name]
    submission.append({
        "sample_id": sample_id,
        "target": pred_value
    })

submission_df = pd.DataFrame(submission)

# Verify submission format
print(f"\n{'='*60}")
print("SUBMISSION VERIFICATION")
print(f"{'='*60}")
print(f"Number of rows: {len(submission_df)}")
print(f"Columns: {submission_df.columns.tolist()}")
print(f"Any NaN values: {submission_df['target'].isna().sum()}")
print(f"Any negative values: {(submission_df['target'] < 0).sum()}")

# Save submission
submission_df.to_csv("submission.csv", index=False)

print(f"\nCreated submission.csv with {len(submission_df)} rows")
print("\nSample submission (first 10 rows):")
print(submission_df.head(10))

print("\nPrediction statistics:")
for target_name in ["Dry_Green_g", "Dry_Dead_g", "Dry_Clover_g", "GDM_g", "Dry_Total_g"]:
    target_preds = submission_df[submission_df["sample_id"].str.contains(target_name)]["target"]
    if len(target_preds) > 0:
        print(f"  {target_name}: mean={target_preds.mean():.2f}, std={target_preds.std():.2f}, min={target_preds.min():.2f}, max={target_preds.max():.2f}")