import logging
import os
from pathlib import Path
from typing import Callable, Dict, List, Optional, Tuple, Union

import h5py
import numpy as np
import torch
import xarray as xr
from PIL import Image
from torch.utils.data import Dataset  # This line imports Dataset from PyTorch
import fsspec
import gcsfs

class WeatherDataset:
    """Dataset class for loading weather data from HDF5 files."""

    def __init__(self, data_path: str, variables: List[str]):
        """Initialize the dataset.
        
        Args:
            data_path: Path to the data directory
            variables: List of variable names to load
        """
        self.data_path = Path(data_path)
        self.variables = variables
        self._load_data()

    def _load_data(self):
        """Load data from HDF5 files."""
        self.data = {}
        for var in self.variables:
            file_path = self.data_path / f"{var}_train.h5"
            if file_path.exists():
                with h5py.File(file_path, "r") as f:
                    self.data[var] = np.array(f[var])
            else:
                print(f"Warning: File {file_path} not found.")

    def __len__(self) -> int:
        """Return the number of samples."""
        if not self.data:
            return 0
        return len(next(iter(self.data.values())))

    def __getitem__(self, idx: int) -> Dict[str, np.ndarray]:
        """Get a sample from the dataset."""
        return {var: data[idx] for var, data in self.data.items()}


class StyleTransferDataset(Dataset):
    """Dataset for paired or unpaired style transfer tasks.

    This dataset is intentionally format-agnostic so the library can ingest
    a wide range of data types (NumPy arrays, PIL images, PyTorch tensors,
    or file paths). It enables replacing pix2pix/CycleGAN-style preprocessing
    with flow matching by returning the content input, a target frame, and an
    optional style reference in a unified dictionary structure.
    """

    def __init__(
        self,
        content_items: List[Union[str, np.ndarray, torch.Tensor, Image.Image]],
        target_items: Optional[
            List[Union[str, np.ndarray, torch.Tensor, Image.Image]]
        ] = None,
        style_items: Optional[
            List[Union[str, np.ndarray, torch.Tensor, Image.Image]]
        ] = None,
        transform: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
        style_transform: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
        use_target_as_style: bool = True,
    ):
        """Create a style transfer dataset.

        Args:
            content_items: Items representing the source domain (paths, arrays,
                tensors, or PIL images).
            target_items: Items representing the target domain. If ``None`` the
                content is reused as the target.
            style_items: Optional explicit style references. If omitted and
                ``use_target_as_style`` is True, the target will be used as the
                style signal.
            transform: Optional transform applied to both content and target
                tensors after conversion.
            style_transform: Optional transform applied only to the style
                tensor after conversion.
            use_target_as_style: When True, fall back to using the target frame
                as the style reference when an explicit ``style_items`` list is
                not provided.
        """

        if target_items is None:
            target_items = content_items

        if len(content_items) != len(target_items):
            raise ValueError("content_items and target_items must be the same length")

        if style_items is not None and len(style_items) != len(content_items):
            raise ValueError("style_items must be None or match the content length")

        self.content_items = content_items
        self.target_items = target_items
        self.style_items = style_items
        self.transform = transform
        self.style_transform = style_transform
        self.use_target_as_style = use_target_as_style

    def __len__(self) -> int:
        return len(self.content_items)

    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
        content = self._to_tensor(self.content_items[idx])
        target = self._to_tensor(self.target_items[idx])

        style_source = None
        if self.style_items is not None:
            style_source = self.style_items[idx]
        elif self.use_target_as_style:
            style_source = self.target_items[idx]

        style = self._to_tensor(style_source) if style_source is not None else content

        if self.transform:
            content = self.transform(content)
            target = self.transform(target)

        if self.style_transform:
            style = self.style_transform(style)
        elif self.transform:
            style = self.transform(style)

        return {
            "input": content,
            "target": target,
            "style": style,
            "metadata": {
                "index": idx,
                "has_explicit_style": self.style_items is not None,
            },
        }

    def _to_tensor(
        self, item: Union[str, Path, np.ndarray, torch.Tensor, Image.Image, None]
    ) -> torch.Tensor:
        """Convert various item types to a float32 CHW tensor in [0, 1]."""

        if item is None:
            raise ValueError("Cannot convert None to tensor for StyleTransferDataset")

        if isinstance(item, torch.Tensor):
            tensor = item.float()
        elif isinstance(item, np.ndarray):
            tensor = torch.from_numpy(item).float()
        elif isinstance(item, Image.Image):
            tensor = torch.from_numpy(np.array(item)).float()
        elif isinstance(item, (str, Path)):
            with Image.open(str(item)) as img:
                tensor = torch.from_numpy(np.array(img)).float()
        else:
            raise TypeError(f"Unsupported item type: {type(item)}")

        if tensor.ndim == 2:
            tensor = tensor.unsqueeze(0)
        elif tensor.ndim == 3:
            # Heuristically detect if data is already channel-first
            if tensor.shape[0] in (1, 3, 4) and tensor.shape[0] <= tensor.shape[-1]:
                pass
            else:
                tensor = tensor.permute(2, 0, 1)
        else:
            raise ValueError(f"Tensor has unsupported shape {tensor.shape}")

        # Normalize if data looks like integer imagery
        if tensor.dtype.is_floating_point:
            if tensor.max() > 1.0:
                tensor = tensor / 255.0
        else:
            tensor = tensor.float() / 255.0

        return tensor

class ERA5Dataset(Dataset):
    """Dataset class for loading ERA5 reanalysis data from WeatherBench 2."""
    
    VARIABLE_MAPPING = {
        't': 'temperature',
        'z': 'geopotential',
        'u': 'u_component_of_wind',
        'v': 'v_component_of_wind'
    }
    
    DEFAULT_URL = "gs://weatherbench2/datasets/era5/1959-2023_01_10-6h-64x32_equiangular_conservative.zarr"
    
    def __init__(
        self,
        data_path: Optional[str] = None,
        resolution: str = "64x32",
        variables: List[str] = ['t', 'z'],
        pressure_levels: List[int] = [850, 700, 500, 300, 200],
        time_slice: Union[slice, str, Tuple[str, str]] = slice('2015', '2016')
    ):
        """Initialize ERA5Dataset."""
        super().__init__()
        
        self.data_path = data_path or self.DEFAULT_URL
        if isinstance(time_slice, tuple):
            time_slice = slice(*time_slice)
            
        self.variables = [self.VARIABLE_MAPPING[v] for v in variables]
        self.pressure_levels = pressure_levels
        
        print(f"Loading data from: {self.data_path}")
        self._load_data(time_slice)
        
    def _load_data(self, time_slice: slice):
        """Load the dataset and select time period."""
        methods = [
            # Method 1: Simple anonymous access
            lambda: {
                'method': xr.open_zarr,
                'args': [self.data_path],
                'kwargs': {
                    'storage_options': {'anon': True},
                    'consolidated': True
                }
            },
            
            # Method 2: Direct HTTP access
            lambda: {
                'method': xr.open_zarr,
                'args': [fsspec.filesystem(
                    'http',
                    client_kwargs={
                        'trust_env': False,
                        'timeout': 30
                    }
                ).get_mapper(self.data_path.replace('gs://', 'https://storage.googleapis.com/'))],
                'kwargs': {'consolidated': True}
            },
            
            # Method 3: GCS anonymous access
            lambda: {
                'method': xr.open_zarr,
                'args': [gcsfs.GCSFileSystem(token='anon').get_mapper(self.data_path)],
                'kwargs': {'consolidated': True}
            },
            
            # Method 4: Configured timeouts
            lambda: {
                'method': xr.open_zarr,
                'args': [self.data_path],
                'kwargs': {
                    'storage_options': {
                        'anon': True,
                        'timeout': 30,
                        'retries': 10,
                        'default_fill_cache': False
                    },
                    'consolidated': True
                }
            }
        ]

        last_exception = None
        for method_factory in methods:
            try:
                method_info = method_factory()
                logging.info(f"Attempting to open dataset with method: {method_info}")
                self.ds = method_info['method'](
                    *method_info['args'],
                    **method_info['kwargs']
                )
                self.times = self.ds.time.sel(time=time_slice)
                print(f"Selected time period: {self.times[0].values} to {self.times[-1].values}")
                print(f"Variables: {self.variables}")
                print(f"Pressure levels: {self.pressure_levels}")
                return  # Success! Exit the method
                
            except Exception as e:
                last_exception = e
                logging.warning(f"Method failed with error: {str(e)}")
                continue
        
        # If we get here, all methods failed
        raise RuntimeError(f"All methods to load data failed. Last error: {str(last_exception)}")

    def __len__(self) -> int:
        """Return number of samples (time steps - 1 for input/target pairs)."""
        return len(self.times) - 1
    
    def __getitem__(self, idx: int) -> Dict[str, Union[torch.Tensor, Dict]]:
        """Get a single sample with current and next timestep."""
        t0 = self.times[idx].values
        t1 = self.times[idx + 1].values
        
        data_t0 = {}
        data_t1 = {}
        
        for var in self.variables:
            data_t0[var] = self.ds[var].sel(
                time=t0,
                level=self.pressure_levels
            ).values
            
            data_t1[var] = self.ds[var].sel(
                time=t1,
                level=self.pressure_levels
            ).values
        
        input_data = torch.tensor(np.stack([data_t0[var] for var in self.variables]))
        target_data = torch.tensor(np.stack([data_t1[var] for var in self.variables]))
        
        return {
            'input': input_data,
            'target': target_data,
            'metadata': {
                't0': t0,
                't1': t1,
                'variables': self.variables,
                'pressure_levels': self.pressure_levels
            }
        }

    @property
    def shape(self) -> Tuple[int, ...]:
        """Return the shape of the data (n_vars, n_levels, lat, lon)."""
        return (len(self.variables), len(self.pressure_levels), 
                self.ds.latitude.size, self.ds.longitude.size)
