from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np

from .sensors import SENSOR_BANDS

@dataclass
class RasterData:
    array: np.ndarray  # shape: (bands, height, width)
    band_names: list[str]
    crs: Any | None
    transform: Any | None
    profile: dict[str, Any]
    nodata: float | None
    source: str
    sensor: str | None = None


def load_tabular(path: Path):
    import pandas as pd

    if path.suffix.lower() in {".csv"}:
        return pd.read_csv(path)
    if path.suffix.lower() in {".xlsx", ".xls"}:
        return pd.read_excel(path)
    raise ValueError("Unsupported tabular format. Use CSV or XLSX.")


def _guess_sensor_from_name(name: str) -> str | None:
    tokens = {
        "sentinel-2": "Sentinel-2",
        "sentinel2": "Sentinel-2",
        "s2a": "Sentinel-2",
        "s2b": "Sentinel-2",
        "s2": "Sentinel-2",
        "sentinel-1": "Sentinel-1",
        "sentinel1": "Sentinel-1",
        "s1": "Sentinel-1",
        "sar": "Sentinel-1",
        "pol": "Sentinel-1",
        "sentinel-3": "Sentinel-3",
        "sentinel3": "Sentinel-3",
        "s3": "Sentinel-3",
        "landsat": "Landsat",
        "planet": "PlanetScope",
        "planetscope": "PlanetScope",
        "uav": "UAV",
        "drone": "UAV",
        "era5": "ERA5",
        "enmap": "ENMAP",
        "hyperspectral": "Hyperspectral",
        "hyper": "Hyperspectral",
    }
    lowered = name.lower()
    for token, sensor in tokens.items():
        if token in lowered:
            return sensor
    return None


def _guess_sensor_from_band_count(bands: int) -> str | None:
    count_map = {
        13: "Sentinel-2",
        2: "Sentinel-1",
        21: "Sentinel-3",
        10: "Landsat",
        4: "PlanetScope",
        5: "ERA5",
        30: "ENMAP",
        50: "Hyperspectral",
    }
    return count_map.get(bands)


def load_npy(path: Path) -> RasterData:
    import numpy as np
    import json

    array = np.load(path, allow_pickle=False)

    def _is_square(n: int) -> tuple[bool, int]:
        root = int(np.sqrt(n))
        return root * root == n, root

    meta_path = path.with_suffix(".json")
    meta: dict[str, Any] = {}
    if meta_path.exists():
        try:
            meta = json.loads(meta_path.read_text(encoding="utf-8"))
        except Exception:
            meta = {}

    sensor_hint = None
    if meta:
        sensor_hint = meta.get("sensor")
    if not sensor_hint:
        sensor_hint = _guess_sensor_from_name(path.stem)

    def _reshape_flat_pixels(arr: np.ndarray) -> tuple[np.ndarray, dict[str, Any]]:
        info: dict[str, Any] = {}
        min_pixels = 1024
        # 2D (pixels, bands) or (bands, pixels)
        if arr.ndim == 2:
            if arr.shape[1] <= 32:
                is_sq, root = _is_square(int(arr.shape[0]))
                if is_sq and arr.shape[0] >= min_pixels:
                    info["reshape"] = "pixels_bands"
                    info["full_shape"] = (root, root)
                    reshaped = arr.reshape(root, root, arr.shape[1]).transpose(2, 0, 1)
                    return reshaped, info
            if arr.shape[0] <= 32:
                is_sq, root = _is_square(int(arr.shape[1]))
                if is_sq and arr.shape[1] >= min_pixels:
                    info["reshape"] = "bands_pixels"
                    info["full_shape"] = (root, root)
                    reshaped = arr.reshape(arr.shape[0], root, root)
                    return reshaped, info
        # 3D (time, pixels, bands) or similar
        if arr.ndim == 3:
            if arr.shape[2] <= 32:
                is_sq, root = _is_square(int(arr.shape[1]))
                if is_sq and arr.shape[1] >= min_pixels:
                    t, pixels, bands = arr.shape
                    info["reshape"] = "time_pixels_bands"
                    info["full_shape"] = (root, root)
                    info["time"] = t
                    info["bands_per_time"] = bands
                    reshaped = arr.reshape(t, root, root, bands).transpose(0, 3, 1, 2)
                    reshaped = reshaped.reshape(t * bands, root, root)
                    return reshaped, info
            if arr.shape[1] <= 32:
                is_sq, root = _is_square(int(arr.shape[0]))
                if is_sq and arr.shape[0] >= min_pixels:
                    pixels, bands, t = arr.shape
                    info["reshape"] = "pixels_bands_time"
                    info["full_shape"] = (root, root)
                    info["time"] = t
                    info["bands_per_time"] = bands
                    reshaped = arr.reshape(root, root, bands, t).transpose(3, 2, 0, 1)
                    reshaped = reshaped.reshape(t * bands, root, root)
                    return reshaped, info
        return arr, info

    def _ensure_bands_first(arr: np.ndarray, hint: str | None) -> np.ndarray:
        if arr.ndim == 2:
            return arr[None, :, :]
        if arr.ndim < 2:
            raise ValueError("Unsupported NPY shape. Expect 2D+ array.")
        if arr.ndim == 3:
            # If a sensor hint is present, align that band count to band axis
            if hint and hint in SENSOR_BANDS:
                target = len(SENSOR_BANDS[hint])
                if target in arr.shape:
                    band_axis = arr.shape.index(target)
                    return np.moveaxis(arr, band_axis, 0)
        # Choose two largest dimensions as spatial axes
        axes = list(range(arr.ndim))
        spatial_axes = sorted(axes, key=lambda i: arr.shape[i], reverse=True)[:2]
        band_axes = [i for i in axes if i not in spatial_axes]
        # Move to (band_axes..., spatial_axes...)
        new_order = band_axes + spatial_axes
        arr = np.moveaxis(arr, new_order, list(range(arr.ndim)))
        bands = int(np.prod([arr.shape[i] for i in range(len(band_axes))])) if band_axes else 1
        h = arr.shape[len(band_axes)]
        w = arr.shape[len(band_axes) + 1]
        return arr.reshape((bands, h, w))

    reshape_info: dict[str, Any] = {}
    if array.ndim >= 2:
        array, reshape_info = _reshape_flat_pixels(array)
        if not reshape_info:
            array = _ensure_bands_first(array, sensor_hint)
    else:
        raise ValueError("Unsupported NPY shape. Expect 2D+ array.")

    band_count = array.shape[0]
    sensor = None
    reason = None

    if meta:
        sensor = meta.get("sensor")
        if sensor:
            reason = "sidecar metadata"

    if not sensor:
        sensor = _guess_sensor_from_name(path.stem)
        if sensor:
            reason = "filename hint"

    if not sensor:
        sensor = _guess_sensor_from_band_count(band_count)
        if sensor:
            reason = "band count"

    band_names = None
    if isinstance(meta.get("band_names"), list) and len(meta["band_names"]) == band_count:
        band_names = list(meta["band_names"])
    elif sensor == "Sentinel-1" and band_count == 4:
        band_names = ["VV", "VH", "HH", "HV"]
    elif sensor in SENSOR_BANDS and len(SENSOR_BANDS[sensor]) == band_count and SENSOR_BANDS[sensor]:
        band_names = list(SENSOR_BANDS[sensor])

    if band_names is None:
        if reshape_info.get("time") and reshape_info.get("bands_per_time"):
            band_names = []
            t = int(reshape_info["time"])
            b = int(reshape_info["bands_per_time"])
            for ti in range(t):
                for bi in range(b):
                    band_names.append(f"T{ti}_B{bi+1}")
        else:
            band_names = [f"band_{idx+1}" for idx in range(band_count)]

    preview = False
    preview_stride = 1
    full_shape = (array.shape[1], array.shape[2])
    bytes_est = array.size * array.dtype.itemsize
    max_bytes = 500_000_000  # ~500MB
    if bytes_est > max_bytes:
        max_side = 1500
        preview_stride = max(1, int(max(full_shape) // max_side))
        if preview_stride > 1:
            array = array[:, ::preview_stride, ::preview_stride]
            preview = True
    return RasterData(
        array=array.astype("float32"),
        band_names=band_names,
        crs=None,
        transform=None,
        profile={
            "sensor_guess": sensor,
            "sensor_guess_reason": reason,
            "preview": preview,
            "preview_stride": preview_stride,
            "full_shape": full_shape,
            **reshape_info,
        },
        nodata=None,
        source=str(path),
        sensor=sensor,
    )

def load_geotiff(path: Path) -> RasterData:
    import rasterio

    with rasterio.open(path) as src:
        array = src.read().astype("float32")
        band_names = [f"band_{idx+1}" for idx in range(array.shape[0])]
        return RasterData(
            array=array,
            band_names=band_names,
            crs=src.crs,
            transform=src.transform,
            profile=src.profile,
            nodata=src.nodata,
            source=str(path),
            sensor=None,
        )


def load_netcdf(path: Path, variable: str | None = None) -> RasterData:
    import xarray as xr

    ds = xr.open_dataset(path)
    data_vars = list(ds.data_vars)
    if not data_vars:
        raise ValueError("No data variables found in NetCDF.")

    arrays = []
    band_names = []

    def _reduce_to_2d(values: np.ndarray) -> np.ndarray:
        arr = values
        while arr.ndim > 2:
            arr = arr[0]
        if arr.ndim != 2:
            raise ValueError("Could not reduce NetCDF variable to 2D.")
        return arr

    if variable is None:
        for var in data_vars:
            data = ds[var]
            values = _reduce_to_2d(data.values)
            arrays.append(values)
            band_names.append(var)
    else:
        data = ds[variable]
        values = data.values
        if values.ndim == 2:
            arrays.append(values)
            band_names.append(variable)
        elif values.ndim >= 3:
            reduced = _reduce_to_2d(values)
            arrays.append(reduced)
            band_names.append(variable)
        else:
            raise ValueError("Unsupported NetCDF variable dimensionality.")

    array = np.stack(arrays, axis=0)
    crs = getattr(data, "rio", None).crs if hasattr(data, "rio") else None
    transform = getattr(data, "rio", None).transform() if hasattr(data, "rio") else None

    return RasterData(
        array=array.astype("float32"),
        band_names=band_names,
        crs=crs,
        transform=transform,
        profile={},
        nodata=None,
        source=str(path),
        sensor="ERA5",
    )


def load_envi(hdr_path: Path) -> RasterData:
    from spectral import open_image

    img = open_image(str(hdr_path))
    array = img.load().astype("float32")
    if array.ndim == 3:
        array = np.moveaxis(array, -1, 0)
    band_names = [f"band_{idx+1}" for idx in range(array.shape[0])]
    return RasterData(
        array=array,
        band_names=band_names,
        crs=None,
        transform=None,
        profile={},
        nodata=None,
        source=str(hdr_path),
        sensor="ENMAP",
    )


def scan_sentinel2_folder(folder: Path) -> dict[str, Path]:
    band_files = {}
    for path in folder.rglob("*.jp2"):
        name = path.name
        if "_B" in name:
            parts = name.split("_B")
            if len(parts) > 1:
                band_code = parts[1].split("_")[0]
                band_files[f"B{band_code}"] = path
    return band_files


def scan_sentinel2_tiffs(folder: Path) -> dict[str, Path]:
    band_files = {}
    for path in folder.rglob("*.tif*"):
        name = path.name
        if "_B" in name:
            parts = name.split("_B")
            if len(parts) > 1:
                band_code = parts[1].split("_")[0].replace("A", "A")
                band_files[f"B{band_code}"] = path
    return band_files


def load_sentinel2_band_tiffs(band_paths: dict[str, Path]) -> RasterData:
    import rasterio
    from rasterio.enums import Resampling

    if not band_paths:
        raise ValueError("No Sentinel-2 band files provided.")

    # choose base with largest pixel count
    def _pixel_count(p: Path) -> int:
        with rasterio.open(p) as ds:
            return ds.width * ds.height

    sorted_bands = sorted(band_paths.items())
    base_band, base_path = max(sorted_bands, key=lambda item: _pixel_count(item[1]))

    arrays = []
    names = []
    crs = None
    transform = None
    profile = {}
    nodata = None

    with rasterio.open(base_path) as base_ds:
        base_h, base_w = base_ds.height, base_ds.width
        crs = base_ds.crs
        transform = base_ds.transform
        profile = base_ds.profile
        nodata = base_ds.nodata

        for band_name, path in sorted_bands:
            with rasterio.open(path) as src:
                data = src.read(
                    1,
                    out_shape=(base_h, base_w),
                    resampling=Resampling.bilinear,
                ).astype("float32")
                arrays.append(data)
                names.append(band_name)

    array = np.stack(arrays, axis=0)
    return RasterData(
        array=array,
        band_names=names,
        crs=crs,
        transform=transform,
        profile=profile,
        nodata=nodata,
        source=str(base_path.parent),
        sensor="Sentinel-2",
    )


def _extract_s2_band_name(path: Path) -> str:
    name = path.stem
    if "_B" in name:
        parts = name.split("_B")
        if len(parts) > 1:
            band = parts[1].split("_")[0]
            return f"B{band}"
    return name


def load_sentinel2_bands(band_paths: list[Path]) -> RasterData:
    import rasterio

    arrays = []
    names = []
    crs = None
    transform = None
    profile = {}
    nodata = None

    for path in band_paths:
        with rasterio.open(path) as src:
            arrays.append(src.read(1).astype("float32"))
            names.append(_extract_s2_band_name(path))
            crs = src.crs
            transform = src.transform
            profile = src.profile
            nodata = src.nodata

    array = np.stack(arrays, axis=0)
    return RasterData(
        array=array,
        band_names=names,
        crs=crs,
        transform=transform,
        profile=profile,
        nodata=nodata,
        source=str(band_paths[0].parent),
        sensor="Sentinel-2",
    )


def save_geotiff(path: Path, array: np.ndarray, reference: RasterData) -> None:
    import rasterio

    bands, height, width = array.shape
    profile = reference.profile.copy()
    profile.update(
        {
            "count": bands,
            "height": height,
            "width": width,
            "dtype": array.dtype,
            "transform": reference.transform,
            "crs": reference.crs,
        }
    )
    path.parent.mkdir(parents=True, exist_ok=True)
    with rasterio.open(path, "w", **profile) as dst:
        dst.write(array)
