import os
import json
import random
import numpy as np
import torch
import geopandas as gpd
from torch.utils.data import Dataset


PASTIS_CLASSES = {
    0:'Background', 1:'Meadow', 2:'Soft winter wheat', 3:'Corn',
    4:'Winter barley', 5:'Winter rapeseed', 6:'Spring barley',
    7:'Sunflower', 8:'Grapevine', 9:'Beet', 10:'Winter triticale',
    11:'Winter durum wheat', 12:'Fruits vegetables flowers',
    13:'Potatoes', 14:'Leguminous fodder', 15:'Soybeans',
    16:'Orchard', 17:'Mixed cereal', 18:'Sorghum', 19:'Void label',
}
NUM_CLASSES  = 20
IGNORE_INDEX = 19


class PASTISDataset(Dataset):
    """
    PASTIS dataset loader using OFFICIAL fold assignments from metadata.geojson.

    5 official folds (~487-496 patches each), geographically stratified.
    Split strategy:
        fold=1: test=1, val=2, train=3+4+5
        fold=2: test=2, val=3, train=4+5+1
        fold=3: test=3, val=4, train=5+1+2
        fold=4: test=4, val=5, train=1+2+3
        fold=5: test=5, val=1, train=2+3+4

    This matches the standard PASTIS benchmark evaluation protocol.
    """

    def __init__(
        self,
        root,
        fold=1,
        split='train',
        num_frames=32,
        label_channel=0,
        norm=True,
        augment=True,
    ):
        super().__init__()
        self.root          = root
        self.fold          = fold
        self.split         = split
        self.num_frames    = num_frames
        self.label_channel = label_channel
        self.norm          = norm
        self.augment       = augment and (split == 'train')

        # Load per-fold normalization stats
        with open(os.path.join(root, 'NORM_S2_patch.json'), 'r') as f:
            norm_data = json.load(f)
        fold_key  = f'Fold_{fold}'
        self.mean = np.array(norm_data[fold_key]['mean'], dtype=np.float32)
        self.std  = np.array(norm_data[fold_key]['std'],  dtype=np.float32)

        # Load official fold assignments from metadata.geojson
        meta     = gpd.read_file(os.path.join(root, 'metadata.geojson'))
        fold_map = {}
        for _, row in meta.iterrows():
            pid  = str(int(row['ID_PATCH']))
            f    = int(row['Fold'])
            fold_map.setdefault(f, []).append(pid)

        # Standard 5-fold split
        all_folds  = [1, 2, 3, 4, 5]
        test_fold  = fold
        val_fold   = (fold % 5) + 1          # 1→2, 2→3, 3→4, 4→5, 5→1
        train_folds= [f for f in all_folds
                      if f != test_fold and f != val_fold]

        if split == 'test':
            raw_ids = fold_map.get(test_fold, [])
        elif split == 'val':
            raw_ids = fold_map.get(val_fold, [])
        else:
            raw_ids = []
            for tf in train_folds:
                raw_ids.extend(fold_map.get(tf, []))

        # Keep only patches that have both S2 and TARGET files
        s2_dir  = os.path.join(root, 'DATA_S2')
        ann_dir = os.path.join(root, 'ANNOTATIONS')
        self.ids = sorted([
            pid for pid in raw_ids
            if (os.path.exists(os.path.join(s2_dir,  f'S2_{pid}.npy')) and
                os.path.exists(os.path.join(ann_dir, f'TARGET_{pid}.npy')))
        ])

        print(f"PASTIS fold={fold} split={split}: {len(self.ids)} patches "
              f"(augment={self.augment})")

    def __len__(self):
        return len(self.ids)

    def _sample_frames(self, data, num_frames):
        """data: (T, C, H, W) → (num_frames, C, H, W)"""
        T = data.shape[0]
        if T <= num_frames:
            pad = num_frames - T
            if pad > 0:
                data = np.concatenate(
                    [data, np.tile(data[[-1]], (pad, 1, 1, 1))], axis=0)
            return data
        indices = np.linspace(0, T - 1, num_frames, dtype=int)
        return data[indices]

    def _normalize(self, x):
        """x: (T, C, H, W) float32"""
        mean = self.mean[None, :, None, None]
        std  = self.std[None,  :, None, None]
        std  = np.where(std < 1e-6, 1.0, std)
        return (x - mean) / std

    def _augment(self, s2, label):
        """
        s2:    (T, C, H, W) float32 numpy
        label: (H, W)       int64   numpy
        """
        # Random horizontal flip
        if random.random() < 0.5:
            s2    = s2[:, :, :, ::-1].copy()
            label = label[:, ::-1].copy()

        # Random vertical flip
        if random.random() < 0.5:
            s2    = s2[:, :, ::-1, :].copy()
            label = label[::-1, :].copy()

        # Random 90/180/270 rotation
        k = random.randint(0, 3)
        if k > 0:
            s2    = np.rot90(s2,    k, axes=(2, 3)).copy()
            label = np.rot90(label, k, axes=(0, 1)).copy()

        return s2, label

    def __getitem__(self, idx):
        pid = self.ids[idx]

        # Load S2
        s2 = np.load(
            os.path.join(self.root, 'DATA_S2', f'S2_{pid}.npy')
        ).astype(np.float32)                          # (61, 10, 128, 128)

        s2 = self._sample_frames(s2, self.num_frames) # (T, 10, 128, 128)

        if self.norm:
            s2 = self._normalize(s2)

        # Load label
        label = np.load(
            os.path.join(self.root, 'ANNOTATIONS', f'TARGET_{pid}.npy')
        )                                             # (3, 128, 128)
        label = label[self.label_channel]             # (128, 128)

        # Augment (train only)
        if self.augment:
            s2, label = self._augment(s2, label)

        s2    = torch.from_numpy(s2.copy()).float()
        label = torch.from_numpy(label.astype(np.int64).copy())

        return {'S2': s2, 'label': label, 'id': pid}