"""
GetTide. Supports FES2014, FES2014_extrapolated, FES2022 and FES2022_extrapolated.

Model-key to subdirectory mapping (explicit):
- FES2014              -> fes2014/ocean_tide
- FES2014_extrapolated -> fes2014/ocean_tide_extrapolated
- FES2022              -> fes2022b/ocean_tide_20241025
- FES2022_extrapolated -> fes2022b/ocean_tide_extrapolated

Notes:
- pyTMD "model keys" may include extrapolated variants depending on pyTMD version.
  If an extrapolated key is not present in pyTMD's database, this module falls back
  to the non-extrapolated key for reading the constituent list, while still reading
  the extrapolated NetCDF files from disk.
- pyTMD FES extractor VERSION must be FES2014 or FES2022. For extrapolated model keys,
  VERSION is derived by stripping the "_extrapolated" suffix.
"""

from datetime import datetime, timedelta
from pathlib import Path
import pickle
import os
from typing import Union, Tuple

import numpy as np
import pandas as pd
import geopandas as gpd
import xarray as xr
from tqdm.auto import tqdm
from shapely import LineString
from pyproj import Geod, CRS

# Import the tidal prediction, pyTMD
import pyTMD
from pyTMD.io.FES import extract_constants as extract_FES_constants
from pyTMD.predict import time_series as predict_tidal_ts
from tfv_get_tools.utilities.parsers import _parse_date, _parse_path
from tfv_get_tools.utilities.warnings import deprecated
from tfv_get_tools.tide._nodestring import (
    load_nodestring_shapefile,
    process_nodestring_gdf,
)
from tfv_get_tools.fvc import write_tide_fvc

crs = CRS.from_epsg(4326)

# Model keys supported by this module (pyTMD database keys)
VALID_MODEL_KEYS = {"FES2014", "FES2014_extrapolated", "FES2022", "FES2022_extrapolated"}
# FES "versions" supported by pyTMD FES extractor
VALID_FES_VERSIONS = {"FES2014", "FES2022"}


def _fes_version_from_model_key(model_key: str) -> str:
    """Map pyTMD model key -> FES extractor VERSION."""
    if model_key not in VALID_MODEL_KEYS:
        raise ValueError(
            f"Requested source {model_key} not supported. "
            f"Valid sources: {VALID_MODEL_KEYS}"
        )
    if model_key.endswith("_extrapolated"):
        return model_key.replace("_extrapolated", "")
    return model_key


def _pytmd_root_from_any_model_path(path: Path) -> Path:
    """
    Resolve the pyTMD root directory expected by pyTMD.io.model(...).

    pyTMD database paths for FES are typically relative to a directory that contains
    'fes2014/' or 'fes2022b/' etc. Therefore, we want the parent directory ABOVE the
    fes folder.
    """
    path = Path(path).resolve()

    # If user passed an ocean_tide-style folder, promote to fes directory first
    if path.name in (
        "ocean_tide",
        "ocean_tide_extrapolated",
        "ocean_tide_20241025",
        "load_tide",
        "load_tide_extrapolated",
    ):
        fes_dir = path.parent
    else:
        fes_dir = path

    # If we are sitting in a known fes directory, root is its parent
    if fes_dir.name.lower() in ("fes2014", "fes2022b", "fes2022"):
        return fes_dir.parent

    # Otherwise: scan upward for a folder that looks like the fes directory
    for parent in [fes_dir] + list(fes_dir.parents):
        if parent.name.lower() in ("fes2014", "fes2022b", "fes2022"):
            return parent.parent

    # Fallback: one level up (best effort)
    return fes_dir.parent


def _select_constituent_dir(model_dir: Path, model_key: str) -> Path:
    """
    Choose the directory that contains the 34 elevation constituent NetCDF files,
    using an explicit model_key -> subdirectory mapping:

    - FES2014              -> ocean_tide
    - FES2014_extrapolated -> ocean_tide_extrapolated
    - FES2022              -> ocean_tide_20241025
    - FES2022_extrapolated -> ocean_tide_extrapolated

    The caller may supply:
    - the constituent folder directly (must match the mapping)
    - the FES folder (fes2014 / fes2022b), in which case the mapped subfolder is used
    - a higher-level root containing the FES folder(s), which will be located
    """
    model_dir = Path(model_dir).resolve()

    subdir_map = {
        "FES2014": "ocean_tide",
        "FES2014_extrapolated": "ocean_tide_extrapolated",
        "FES2022": "ocean_tide_20241025",
        "FES2022_extrapolated": "ocean_tide_extrapolated",
    }

    wanted = subdir_map[model_key].lower()
    here = model_dir.name.lower()

    # If user passed the constituent directory directly, accept only if it matches
    if here == wanted:
        return model_dir

    # If user passed the FES directory, append wanted subfolder
    if here in ("fes2014", "fes2022b", "fes2022"):
        candidate = (model_dir / subdir_map[model_key]).resolve()
        if candidate.exists():
            return candidate
        raise FileNotFoundError(
            f"Expected subfolder '{subdir_map[model_key]}' was not found under {model_dir.as_posix()}"
        )

    # Otherwise treat as a higher-level root and locate the correct FES folder
    fes_folder = "fes2022b" if model_key.startswith("FES2022") else "fes2014"
    fes_candidates = list(model_dir.rglob(fes_folder))
    if not fes_candidates:
        raise FileNotFoundError(
            f"Could not locate '{fes_folder}' under {model_dir.as_posix()}"
        )

    fes_dir = fes_candidates[0].resolve()
    candidate = (fes_dir / subdir_map[model_key]).resolve()
    if candidate.exists():
        return candidate

    raise FileNotFoundError(
        f"Expected subfolder '{subdir_map[model_key]}' was not found under {fes_dir.as_posix()}"
    )


class TidalExtractor:
    """Wrapper for PyTMD operations to enable testing."""

    def extract_fes_constants(self, coords, files, fes_version, interpolate_method):
        """Extract FES constants - wrapped for testing."""
        return extract_FES_constants(
            coords[:, 0],
            coords[:, 1],
            files,
            TYPE="z",
            VERSION=fes_version,
            METHOD=interpolate_method,
            GZIP=False,
            SCALE=1.0 / 100.0,
        )

    def predict_tidal_timeseries(self, tvec, hc, cons):
        """Predict tidal timeseries - wrapped for testing."""
        return predict_tidal_ts(tvec, hc, cons, corrections="FES")


# Default extractor instance
_default_extractor = TidalExtractor()


def _detect_tide_model_source(model_dir: Path):
    """Detect tidal model key based on model_dir."""
    original_model_dir = Path(model_dir)
    model_dir = original_model_dir.resolve()

    search_names = {
        model_dir.name.lower(),
        model_dir.parent.name.lower(),
        str(model_dir).lower(),
    }

    if any("fes2014" in s for s in search_names):
        if any("extrap" in s for s in search_names):
            return "FES2014_extrapolated", original_model_dir
        return "FES2014", original_model_dir

    if any("fes2022" in s for s in search_names):
        if any("extrap" in s for s in search_names):
            return "FES2022_extrapolated", original_model_dir
        return "FES2022", original_model_dir

    return None, original_model_dir


def _get_model_dir(fes_version="FES2014", model_dir: Union[str, Path] = None) -> Path:
    """Get and validate model directory."""
    if fes_version not in VALID_FES_VERSIONS:
        raise ValueError(
            f"Requested FES version {fes_version} not supported. "
            f"Valid versions: {VALID_FES_VERSIONS}"
        )

    if model_dir is None:
        env = f"{fes_version}_DIR"
        if env not in os.environ:
            raise ValueError(
                f"The {env} root directory needs to be supplied, either as "
                f"`model_dir` variable, or as environment variable '{env}'"
            )
        model_dir = os.environ[env]

    model_dir = Path(model_dir)

    if not model_dir.exists():
        raise FileNotFoundError(
            f"{fes_version} model directory ({model_dir.as_posix()}) does not exist"
        )

    return model_dir


def _get_chainage_array(array: np.ndarray):
    """Calculate chainage from coordinates."""
    geod = Geod(ellps="WGS84")
    numCoords = array.shape[0] - 1
    geo = LineString(array)

    stf = 0
    for i in range(0, numCoords):
        point1 = geo.coords[i]
        point2 = geo.coords[i + 1]
        _, _, dist = geod.inv(point1[0], point1[1], point2[0], point2[1])
        stf += dist

    nx = len(geo.xy[0])
    incr = stf / (nx - 1)
    chainage = [incr * x for x in range(1, nx)]
    chainage.insert(0, 0)

    return chainage, nx


def _check_coords(coords):
    """Validate coordinate format."""
    if not isinstance(coords, np.ndarray):
        coords = np.asarray(coords)

        if len(coords.shape) == 1:
            coords = coords.reshape([1, 2])

    if coords.shape[1] != 2:
        raise ValueError(
            "Coordinates should be Nx2 format or tuple/list for single location (x, y)"
        )

    return coords


def _normalise_coordinates(coordinates: Union[tuple, np.ndarray, dict]) -> dict:
    """Convert coordinates to consistent dict format."""
    if isinstance(coordinates, tuple):
        if len(coordinates) != 2:
            raise ValueError("Tuple coordinates must be (lon, lat)")
        return {1: np.asarray(coordinates)[None, :]}

    elif isinstance(coordinates, np.ndarray):
        if coordinates.shape[1] != 2:
            raise ValueError("Array coordinates must be Nx2")
        return {1: coordinates}

    elif isinstance(coordinates, dict):
        if len(coordinates) == 0:
            raise ValueError("No coordinates provided")
        return coordinates

    else:
        raise ValueError("Unsupported coordinate format")


def get_constituents(
    coordinates: Union[tuple, np.ndarray, dict],
    model_dir: Union[str, Path] = None,
    interpolate_method="spline",
    save_cons: Union[str, Path] = None,
    source="FES2014",
    extractor=None,
):
    """Get tidal constituents."""
    if extractor is None:
        extractor = _default_extractor

    # source is a pyTMD model key (e.g. FES2022_extrapolated)
    model_key = source
    fes_version = _fes_version_from_model_key(model_key)

    coordinates = _normalise_coordinates(coordinates)

    cons_dict = dict()
    for bnd_id, coords in coordinates.items():
        coords = _check_coords(coords)

        if coords.shape[0] > 1:
            chainage, nx = _get_chainage_array(coords)
        else:
            chainage = 0
            nx = 1

        # Validate the provided model_dir exists (based on FES VERSION, not model key)
        raw_model_dir = _get_model_dir(fes_version, model_dir)

        # Choose the directory containing elevation constituents
        constituent_dir = _select_constituent_dir(raw_model_dir, model_key)

        # Resolve pyTMD root for database-relative paths
        srcdir = _pytmd_root_from_any_model_path(constituent_dir)

        # List constituent files. Exclude land/mask/grid helper files.
        files = []
        for x in constituent_dir.glob("*.nc"):
            stem = x.stem.lower()
            if any(k in stem for k in ("mask", "land", "grid")):
                continue
            files.append(x)

        file_cons = [f.stem for f in files]
        files = [files[i] for i in np.argsort(file_cons)]

        try:
            model = pyTMD.io.model(srcdir).elevation(model_key)
        except Exception:
            # Fallback if pyTMD database does not contain the extrapolated key
            model = pyTMD.io.model(srcdir).elevation(fes_version)
        cons = model.constituents

        # Expect 34 constituents for elevation for these FES solutions
        if fes_version in ("FES2014", "FES2022"):
            if len(files) != 34:
                raise ValueError(
                    f"Cannot find 34 .nc files for {model_key} "
                    f"(found {len(files)}) in {constituent_dir.as_posix()}"
                )

        print("... extracting constituents from database")
        amp, ph = extractor.extract_fes_constants(
            coords, files, fes_version, interpolate_method
        )

        cons_dict[bnd_id] = dict(
            cons=(amp, ph, cons), geo=(coords, chainage, nx), source=model_key
        )

    if save_cons:
        with open(save_cons, "wb") as f:
            pickle.dump(cons_dict, f)

    return cons_dict


def predict_waterlevel_timeseries(
    time_start: Union[str, pd.Timestamp, datetime],
    time_end: Union[str, pd.Timestamp, datetime],
    freq: Union[str, pd.Timedelta, timedelta] = "15min",
    coords: Union[tuple, np.ndarray, dict] = None,
    source: str = "FES2014",
    model_dir: Union[str, Path] = None,
    interpolate_method: str = "spline",
    constituents: Union[dict, str, Path] = None,
    extractor=None,
):
    """Extract tidal waterlevels for coordinates."""
    if extractor is None:
        extractor = _default_extractor

    timevec = pd.date_range(
        start=pd.Timestamp(time_start),
        end=pd.Timestamp(time_end),
        freq=freq,
    )

    ref_date = pd.Timestamp(1992, 1, 1)
    tvec = (timevec - ref_date).to_numpy().astype(float) / (10**9 * 60 * 60 * 24)

    # Load constituents if needed
    if constituents is None:
        if coords is None:
            raise ValueError("Either coords or constituents must be provided")
        constituents = get_constituents(
            coords,
            model_dir,
            source=source,
            interpolate_method=interpolate_method,
            extractor=extractor,
        )
    else:
        if not isinstance(constituents, dict):
            constituents = Path(constituents)
            if not constituents.exists():
                raise FileNotFoundError(f"Constituents file not found: {constituents}")

            with open(constituents, "rb") as f:
                constituents = pickle.load(f)

    dsset = {}
    for label, dat in constituents.items():
        (amp, ph, cons) = dat["cons"]
        source = dat["source"]
        nx = amp.shape[0]

        # Calculate complex phase in radians
        cph = -1j * ph * np.pi / 180.0
        hc = amp * np.exp(cph)

        print("...Expanding timeseries")
        ha = np.ma.zeros((tvec.shape[0], nx))
        for j in tqdm(range(nx)):
            ha[:, j] = extractor.predict_tidal_timeseries(
                tvec, hc[j].reshape(1, len(cons)), cons
            )

        # Remove invalid values by using nearest valid
        if hasattr(ha, "mask") and ha.mask.any():
            mask = ha.mask[0, :]
            real = np.where(mask == False)[0]
            for idx, cond in enumerate(mask.tolist()):
                if cond:
                    nearest_idx = real[np.argmin((real - idx) ** 2)]
                    ha[:, idx] = ha[:, nearest_idx]

        # Create dataset
        if isinstance(dat["geo"][1], (int, float)):
            chain = np.array([dat["geo"][1]])
            squeeze = True
        else:
            chain = dat["geo"][1]
            squeeze = False

        ds = xr.Dataset(
            coords=dict(time=timevec, chainage=chain),
            data_vars=dict(
                wl=(("time", "chainage"), ha),
                longitude=(("chainage",), dat["geo"][0][:, 0]),
                latitude=(("chainage",), dat["geo"][0][:, 1]),
            ),
        )
        ds.attrs["name"] = label
        ds.attrs["source"] = source
        ds["chainage"].attrs = {"long_name": "projected chainage", "units": "m"}
        ds["wl"].attrs = {"long_name": "tidal waterlevel", "units": "m"}
        ds["longitude"].attrs = crs.cs_to_cf()[1]
        ds["latitude"].attrs = crs.cs_to_cf()[0]

        if squeeze:
            ds = ds.isel(chainage=0)

        dsset[label] = ds

    if len(dsset) == 1:
        return list(dsset.values())[0]
    else:
        return dsset


def ExtractTide(
    time_start: Union[str, pd.Timestamp, datetime],
    time_end: Union[str, pd.Timestamp, datetime],
    fname: Union[str, Path] = None,
    model_dir: Union[str, Path] = None,
    shapefile: Union[str, Path] = None,
    out_path: Path = Path("."),
    process_ids: Union[tuple, list] = None,
    freq: Union[str, pd.Timedelta, timedelta] = "15min",
    spacing: int = 2500,
    attrs=dict(),
    interpolate_method="spline",
    source=None,
    local_tz: Tuple[float, str] = None,
    constituents: Union[str, Path, dict] = None,
    write_netcdf: bool = True,
    write_fvc: bool = True,
    fvc_path: Path = None,
    nc_path_str: str = None,
    extractor=None,
):
    """Full workflow for tidal waterlevel extraction."""
    if extractor is None:
        extractor = _default_extractor

    out_path = _parse_path(out_path)
    time_start = _parse_date(time_start)
    time_end = _parse_date(time_end)

    # Extract constituents if not provided
    if constituents is None:
        if source is None:
            if model_dir is None:
                raise ValueError("Either source or model_dir must be provided")
            source, model_dir = _detect_tide_model_source(Path(model_dir))
            if source is None:
                raise ValueError(
                    "Could not detect tidal model source from model directory"
                )

        if source not in VALID_MODEL_KEYS:
            raise ValueError(
                f"Requested source {source} not supported. "
                f"Valid sources: {VALID_MODEL_KEYS}"
            )

        if shapefile is None:
            raise ValueError("Shapefile required when constituents not provided")

        shapefile = Path(shapefile)
        if not shapefile.exists():
            raise FileNotFoundError(f"Shapefile not found: {shapefile}")

        gdf = load_nodestring_shapefile(shapefile, process_ids=process_ids)
        ns_dat = process_nodestring_gdf(gdf, spacing=spacing)

        print("Running GetTide")
        print("--------------------")
        print("Confirming Request:")
        print(f"...Time Start: {time_start.strftime('%Y-%m-%d %H:%M')}")
        print(f"...Time End: {time_end.strftime('%Y-%m-%d %H:%M')}")
        print(f"...Model Dir: {Path(model_dir).absolute().as_posix()}")
        print(f"...Tidal Data Source: {source}")
        print(f"...Nodestring Name: {shapefile.name}")
        print(f"...Nodestring Folder: {shapefile.parent.absolute().as_posix()}")
        print(f"...Nodestring ID's to Process: {list(ns_dat.keys())}")
        print(f"...Nodestring CRS (EPSG): {gdf.crs.to_epsg()}")
        print(f"...Nodestring Spacing: {spacing:0.1f}m")

        constituents = get_constituents(
            ns_dat,
            model_dir,
            source=source,
            interpolate_method=interpolate_method,
            extractor=extractor,
        )
    else:
        if isinstance(constituents, (str, Path)):
            with open(constituents, "rb") as f:
                constituents = pickle.load(f)
        print(f"...Using pre-extracted constituents file")
        # Get source from constituents
        k = list(constituents.keys())[0]
        source = constituents[k]["source"]

    if local_tz is None:
        print(f"...Timezone: UTC (GMT+0.0)")
    else:
        print(f"...Timezone: UTC (GMT+0.0) AND")
        print(f"...Local Timezone: {local_tz[0]:0.1f}, {local_tz[1]}")

    ns_wlev = predict_waterlevel_timeseries(
        time_start, time_end, freq=freq, constituents=constituents, extractor=extractor
    )

    # Convert to dict if single dataset
    if isinstance(ns_wlev, xr.Dataset):
        ns_wlev = {ns_wlev.attrs["name"]: ns_wlev}

    # Set filename
    if fname is None:
        tsstr = time_start.strftime("%Y%m%d")
        testr = time_end.strftime("%Y%m%d")

        fname = f"{source.upper()}_TIDE_{tsstr}_{testr}.nc"

        if local_tz is not None:
            tzlbl = local_tz[1]
            fname = fname.replace(".", f"_{tzlbl}.")

    outname = out_path / fname

    # Write netcdf
    print(f"Writing dataset: {fname}")
    nc = _netcdf_writer(
        constituents,
        ns_wlev,
        outname,
        time_start,
        time_end,
        freq,
        source,
        write_netcdf=write_netcdf,
        local_tz=local_tz,
        attrs=attrs,
    )

    # TODO: This is hardcoded, but we only have FES support so it's ok for now.
    # We must change this if we add other tidal models.
    info_url = "https://www.aviso.altimetry.fr/en/data/products/auxiliary-products/global-tide-fes.html"

    # Write FVC control file
    if write_fvc:
        nc_path_str = outname.as_posix() if not nc_path_str else nc_path_str
        fvc_fname = fname.replace(".nc", ".fvc")
        fvc_path = out_path if not fvc_path else fvc_path

        write_tide_fvc(
            nc,
            nc_path=nc_path_str,
            output_path=fvc_path,
            filename=fvc_fname,
            source=source,
            info_url=info_url,
        )

    return nc


def _netcdf_writer(
    constituents: dict,
    ns_wlev: dict,
    outname: Path,
    time_start: pd.Timestamp,
    time_end: pd.Timestamp,
    freq: str,
    source: str,
    local_tz=None,
    attrs=dict(),
    write_netcdf=True,
):
    """Write netcdf file."""
    encoding = dict()

    timevec = pd.date_range(
        start=pd.Timestamp(time_start),
        end=pd.Timestamp(time_end),
        freq=freq,
    )

    nc = xr.Dataset(
        coords=dict(time=timevec), attrs=attrs.copy()  # Copy to avoid mutable default
    )
    nc.attrs["source"] = source
    nc["time"].attrs["tz"] = "UTC"
    encoding["time"] = dict(dtype=np.float64, units="days since 1990-01-01 00:00:00")

    # Add local time if requested
    if local_tz is not None:
        tz_offset = local_tz[0]
        tz_name = local_tz[1]
        if tz_name is None:
            tz_name = f"GMT{tz_offset:+}"

        nc["local_time"] = (("time",), timevec + pd.Timedelta(tz_offset, unit="h"))
        nc["local_time"].attrs = dict(tz=tz_name)
        encoding["local_time"] = dict(
            dtype=np.float64, units="days since 1990-01-01 00:00:00"
        )

    for bnd_id in ns_wlev.keys():
        coords, chainage, nx = constituents[bnd_id]["geo"]
        wl_data = ns_wlev[bnd_id]["wl"]

        dimstr = f"ns{bnd_id}_chain"
        chnstr = f"ns{bnd_id}_chainage"
        varstr = f"ns{bnd_id}_wl"

        # Ensure chainage is always an array, even for single points
        chainage_array = np.asarray(chainage).astype(np.float32)
        if chainage_array.ndim == 0:  # scalar case
            chainage_array = np.array([chainage_array])

        nc[chnstr] = ((dimstr,), chainage_array)
        nc[chnstr].attrs["units"] = "m"
        nc[chnstr].attrs["longname"] = f"nodestring {bnd_id} chainage"

        # Handle water level data dimensions properly
        if wl_data.ndim == 1:  # Single point case (squeezed)
            # Reshape to (time, 1) for single chain point
            wl_values = wl_data.values.reshape(-1, 1).astype(np.float32)
        else:  # Multiple points case
            wl_values = wl_data.values.astype(np.float32)

        nc[varstr] = (("time", dimstr), wl_values)
        nc[varstr].attrs["units"] = "m"
        nc[varstr].attrs["long_name"] = f"nodestring {bnd_id} waterlevel"

        nc[f"ns{bnd_id}_longitude"] = ((dimstr,), coords[:, 0], crs.cs_to_cf()[1])
        nc[f"ns{bnd_id}_latitude"] = ((dimstr,), coords[:, 1], crs.cs_to_cf()[0])
        nc[f"ns{bnd_id}_longitude"].attrs["description"] = "X extraction coordinate"
        nc[f"ns{bnd_id}_latitude"].attrs["description"] = "Y extraction coordinate"

        encoding[chnstr] = dict(zlib=True, complevel=1, dtype=np.float32)
        encoding[varstr] = dict(zlib=True, complevel=1, dtype=np.float32)

    if write_netcdf:
        nc.to_netcdf(outname, encoding=encoding)

    return nc


# Keep the deprecated function for backwards compatibility
@deprecated(ExtractTide)
def gen_tfv_tide_netcdf(*args, **kwargs):
    """Deprecated - use ExtractTide instead."""
    # Convert old parameter names to new ones
    if "outname" in kwargs:
        kwargs["fname"] = Path(kwargs.pop("outname")).name
        kwargs["out_path"] = Path(kwargs["fname"]).parent
    if "version" in kwargs:
        kwargs["source"] = kwargs.pop("version")
    if "cons_file" in kwargs:
        kwargs["constituents"] = kwargs.pop("cons_file")

    return ExtractTide(*args, **kwargs)
