from __future__ import annotations

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

import numpy as np

from .sensors import SENSOR_ALIASES


@dataclass(frozen=True)
class IndexDefinition:
    name: str
    formula: str
    sensors: set[str]
    required: list[str]
    description: str = ""


DEFAULT_INDICES: list[IndexDefinition] = [
    IndexDefinition(
        name="NDVI",
        formula="(NIR - RED) / (NIR + RED)",
        sensors={"Sentinel-2", "Landsat", "PlanetScope", "UAV", "Sentinel-3", "ENMAP", "Hyperspectral"},
        required=["NIR", "RED"],
        description="Normalized Difference Vegetation Index",
    ),
    IndexDefinition(
        name="NDWI",
        formula="(GREEN - NIR) / (GREEN + NIR)",
        sensors={"Sentinel-2", "Landsat", "PlanetScope", "UAV", "Sentinel-3"},
        required=["GREEN", "NIR"],
        description="Normalized Difference Water Index",
    ),
    IndexDefinition(
        name="MNDWI",
        formula="(GREEN - SWIR1) / (GREEN + SWIR1)",
        sensors={"Sentinel-2", "Landsat"},
        required=["GREEN", "SWIR1"],
        description="Modified Normalized Difference Water Index",
    ),
    IndexDefinition(
        name="NDBI",
        formula="(SWIR1 - NIR) / (SWIR1 + NIR)",
        sensors={"Sentinel-2", "Landsat"},
        required=["SWIR1", "NIR"],
        description="Normalized Difference Built-up Index",
    ),
    IndexDefinition(
        name="SAVI",
        formula="1.5 * (NIR - RED) / (NIR + RED + 0.5)",
        sensors={"Sentinel-2", "Landsat", "PlanetScope", "UAV"},
        required=["NIR", "RED"],
        description="Soil Adjusted Vegetation Index",
    ),
    IndexDefinition(
        name="EVI",
        formula="2.5 * (NIR - RED) / (NIR + 6 * RED - 7.5 * BLUE + 1)",
        sensors={"Sentinel-2", "Landsat", "PlanetScope", "UAV"},
        required=["NIR", "RED", "BLUE"],
        description="Enhanced Vegetation Index",
    ),
    IndexDefinition(
        name="NDMI",
        formula="(NIR - SWIR1) / (NIR + SWIR1)",
        sensors={"Sentinel-2", "Landsat"},
        required=["NIR", "SWIR1"],
        description="Normalized Difference Moisture Index",
    ),
    IndexDefinition(
        name="GNDVI",
        formula="(NIR - GREEN) / (NIR + GREEN)",
        sensors={"Sentinel-2", "Landsat", "PlanetScope", "UAV"},
        required=["NIR", "GREEN"],
        description="Green Normalized Difference Vegetation Index",
    ),
    IndexDefinition(
        name="NDRE",
        formula="(NIR - RE1) / (NIR + RE1)",
        sensors={"Sentinel-2", "ENMAP", "Hyperspectral"},
        required=["NIR", "RE1"],
        description="Normalized Difference Red Edge Index",
    ),
    IndexDefinition(
        name="MSR",
        formula="(NIR / (RED + 1e-6) - 1) / (sqrt(NIR / (RED + 1e-6)) + 1)",
        sensors={"Sentinel-2", "Landsat", "PlanetScope", "UAV"},
        required=["NIR", "RED"],
        description="Modified Simple Ratio",
    ),
    IndexDefinition(
        name="OSAVI",
        formula="1.16 * (NIR - RED) / (NIR + RED + 0.16)",
        sensors={"Sentinel-2", "Landsat", "PlanetScope", "UAV"},
        required=["NIR", "RED"],
        description="Optimized Soil Adjusted Vegetation Index",
    ),
    IndexDefinition(
        name="EVI2",
        formula="2.5 * (NIR - RED) / (NIR + 2.4 * RED + 1)",
        sensors={"Sentinel-2", "Landsat", "PlanetScope", "UAV"},
        required=["NIR", "RED"],
        description="Two-band Enhanced Vegetation Index",
    ),
    IndexDefinition(
        name="ARVI",
        formula="(NIR - (2 * RED - BLUE)) / (NIR + (2 * RED - BLUE))",
        sensors={"Sentinel-2", "Landsat", "PlanetScope", "UAV"},
        required=["NIR", "RED", "BLUE"],
        description="Atmospherically Resistant Vegetation Index",
    ),
    IndexDefinition(
        name="VARI",
        formula="(GREEN - RED) / (GREEN + RED - BLUE)",
        sensors={"Sentinel-2", "Landsat", "PlanetScope", "UAV"},
        required=["GREEN", "RED", "BLUE"],
        description="Visible Atmospherically Resistant Index",
    ),
    IndexDefinition(
        name="BSI",
        formula="(SWIR1 + RED - NIR - BLUE) / (SWIR1 + RED + NIR + BLUE)",
        sensors={"Sentinel-2", "Landsat"},
        required=["SWIR1", "RED", "NIR", "BLUE"],
        description="Bare Soil Index",
    ),
    IndexDefinition(
        name="NBR",
        formula="(NIR - SWIR2) / (NIR + SWIR2)",
        sensors={"Sentinel-2", "Landsat"},
        required=["NIR", "SWIR2"],
        description="Normalized Burn Ratio",
    ),
    IndexDefinition(
        name="NBR2",
        formula="(SWIR1 - SWIR2) / (SWIR1 + SWIR2)",
        sensors={"Sentinel-2", "Landsat"},
        required=["SWIR1", "SWIR2"],
        description="Normalized Burn Ratio 2",
    ),
    IndexDefinition(
        name="NDSI",
        formula="(GREEN - SWIR1) / (GREEN + SWIR1)",
        sensors={"Sentinel-2", "Landsat"},
        required=["GREEN", "SWIR1"],
        description="Normalized Difference Snow Index",
    ),
    IndexDefinition(
        name="NDGI",
        formula="(GREEN - RED) / (GREEN + RED)",
        sensors={"Sentinel-2", "Landsat", "PlanetScope", "UAV"},
        required=["GREEN", "RED"],
        description="Normalized Difference Greenness Index",
    ),
    IndexDefinition(
        name="MSI",
        formula="SWIR1 / (NIR + 1e-6)",
        sensors={"Sentinel-2", "Landsat"},
        required=["SWIR1", "NIR"],
        description="Moisture Stress Index",
    ),
    IndexDefinition(
        name="SIPI",
        formula="(NIR - BLUE) / (NIR - RED + 1e-6)",
        sensors={"Sentinel-2", "Landsat", "PlanetScope", "UAV"},
        required=["NIR", "BLUE", "RED"],
        description="Structure Insensitive Pigment Index",
    ),
    IndexDefinition(
        name="CIgreen",
        formula="NIR / (GREEN + 1e-6) - 1",
        sensors={"Sentinel-2", "Landsat", "PlanetScope", "UAV"},
        required=["NIR", "GREEN"],
        description="Chlorophyll Index Green",
    ),
    IndexDefinition(
        name="CIre",
        formula="NIR / (RE3 + 1e-6) - 1",
        sensors={"Sentinel-2", "ENMAP", "Hyperspectral"},
        required=["NIR", "RE3"],
        description="Chlorophyll Index Red Edge",
    ),
    IndexDefinition(
        name="VV/VH Ratio",
        formula="VV / (VH + 1e-6)",
        sensors={"Sentinel-1"},
        required=["VV", "VH"],
        description="SAR ratio index",
    ),
    IndexDefinition(
        name="RVI",
        formula="4 * VH / (VV + VH + 1e-6)",
        sensors={"Sentinel-1"},
        required=["VV", "VH"],
        description="Radar Vegetation Index",
    ),
    IndexDefinition(
        name="VV+VH",
        formula="VV + VH",
        sensors={"Sentinel-1"},
        required=["VV", "VH"],
        description="SAR sum",
    ),
    IndexDefinition(
        name="VV-VH",
        formula="VV - VH",
        sensors={"Sentinel-1"},
        required=["VV", "VH"],
        description="SAR difference",
    ),
    IndexDefinition(
        name="VH/VV",
        formula="VH / (VV + 1e-6)",
        sensors={"Sentinel-1"},
        required=["VV", "VH"],
        description="SAR inverse ratio",
    ),
    IndexDefinition(
        name="WindSpeed",
        formula="sqrt(U10 * U10 + V10 * V10)",
        sensors={"ERA5"},
        required=["U10", "V10"],
        description="10m wind speed",
    ),
]


_ALLOWED_FUNCS = {
    "sqrt": np.sqrt,
    "log": np.log,
    "exp": np.exp,
    "abs": np.abs,
    "clip": np.clip,
}


class FormulaError(Exception):
    pass


def get_indices(sensor: str) -> list[IndexDefinition]:
    return [idx for idx in DEFAULT_INDICES if sensor in idx.sensors]


def load_custom_indices(path: Path) -> list[IndexDefinition]:
    if not path.exists():
        raise FileNotFoundError(f"Index file not found: {path}")
    text = path.read_text(encoding="utf-8")
    if path.suffix.lower() in {".yaml", ".yml"}:
        try:
            import yaml  # type: ignore
        except Exception as exc:
            raise ImportError("PyYAML is required for YAML index files.") from exc
        raw = yaml.safe_load(text) or []
    else:
        raw = json.loads(text)
    indices: list[IndexDefinition] = []
    for item in raw:
        indices.append(
            IndexDefinition(
                name=item["name"],
                formula=item["formula"],
                sensors=set(item.get("sensors", ["Custom"])),
                required=item.get("required", []),
                description=item.get("description", ""),
            )
        )
    return indices


def resolve_alias(sensor: str, alias: str) -> str:
    mapping = SENSOR_ALIASES.get(sensor, {})
    return mapping.get(alias, alias)


def safe_eval(formula: str, variables: dict[str, np.ndarray]) -> np.ndarray:
    node = ast.parse(formula, mode="eval")

    def _eval(expr: ast.AST) -> np.ndarray:
        if isinstance(expr, ast.Expression):
            return _eval(expr.body)
        if isinstance(expr, ast.BinOp):
            left = _eval(expr.left)
            right = _eval(expr.right)
            if isinstance(expr.op, ast.Add):
                return left + right
            if isinstance(expr.op, ast.Sub):
                return left - right
            if isinstance(expr.op, ast.Mult):
                return left * right
            if isinstance(expr.op, ast.Div):
                return left / (right + 1e-9)
            if isinstance(expr.op, ast.Pow):
                return left ** right
            raise FormulaError("Unsupported binary operator")
        if isinstance(expr, ast.UnaryOp):
            operand = _eval(expr.operand)
            if isinstance(expr.op, ast.UAdd):
                return operand
            if isinstance(expr.op, ast.USub):
                return -operand
            raise FormulaError("Unsupported unary operator")
        if isinstance(expr, ast.Name):
            if expr.id in variables:
                return variables[expr.id]
            raise FormulaError(f"Unknown variable: {expr.id}")
        if isinstance(expr, ast.Constant):
            return np.array(expr.value, dtype="float32")
        if isinstance(expr, ast.Call):
            if not isinstance(expr.func, ast.Name):
                raise FormulaError("Only simple functions are allowed")
            func_name = expr.func.id
            if func_name not in _ALLOWED_FUNCS:
                raise FormulaError(f"Function {func_name} not allowed")
            args = [_eval(arg) for arg in expr.args]
            return _ALLOWED_FUNCS[func_name](*args)
        raise FormulaError("Unsupported expression")

    return _eval(node)


def compute_index(formula: str, band_map: dict[str, np.ndarray]) -> np.ndarray:
    return safe_eval(formula, band_map)
