# -*- coding: utf-8 -*-
"""
Created on Mon Sep  8 19:42:45 2025

@author: anba
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import ticker as mticker
from matplotlib.colors import LogNorm, Normalize
from matplotlib.ticker import FixedLocator
import cmocean as cmo
from typing import Tuple, List, Optional, Sequence

from .utils_dfsu2d import DfsuUtils2D
from .utils_xml import XMLUtils
from .adcp import ADCP as ADCPDataset
from .utils_crs import CRSHelper
from .plotting import PlottingShell
from .utils_shapefile import ShapefileLayer


def mt_model_transect_comparison(
    mt_model: DfsuUtils2D,
    adcp: ADCPDataset,
    crs_helper: CRSHelper,
    shapefile_layers: Optional[Sequence[ShapefileLayer]] = None,

    # --- ADCP ---
    adcp_series_mode: str = "bin",                # "bin" | "depth" | "hab"
    adcp_series_target: str = "mean",             # "mean" | "pXX"
    adcp_transect_color: str = "magenta",
    

    # --- SSC ---
    ssc_item_number: int = 1,
    ssc_scale: str = "log",                           # "log" or "normal"
    ssc_levels: list[float] = (0.01, 0.1, 1.0, 10.0, 100.0),
    ssc_vmin: float | None = None,                    # mg/L
    ssc_vmax: float | None = None,                    # mg/L
    ssc_cmap_name: str = "turbo",
    ssc_bottom_thresh: float = 0.01,
    ssc_pixel_size_m: float = 10.0,                   # Raster resolution
    
    # --- Layout ---
    cbar_tick_decimals: int = 2,
    axis_tick_decimals: int = 3,
    pad_m: float = 2000.0,
    distance_bin_m: float = 50.0,
    bar_width_scale: float = 0.15,
    
):
    """
    Plot MT model vs. ADCP transect comparison with three stacked subplots:
      1) Top: Map of model SSC raster with ADCP track and optional shapefile overlays.
      2) Middle: Distance-binned bar chart (Model vs ADCP).
      3) Bottom: Metadata columns.

    Parameters
    ----------
    mt_model : DfsuUtils2D
        Preloaded model object. Must provide:
        - rasterize_idw_bbox(item_number, bbox, t, pixel_size_m) -> (array, extent)
        - extract_transect(xq, yq, t, item_number) -> (time_series, ...)
    adcp : ADCPDataset
        Preloaded ADCP object. Must provide:
        - position.x, position.y, position.distance, position.epsg
        - time.ensemble_datetimes
        - get_beam_series(field_name, mode, target) -> (series, meta)
        - name
    crs_helper : CRSHelper
        Must provide:
        - bbox_from_coords(x, y, pad_m, from_crs) -> [xmin, xmax, ymin, ymax]
        - axis_labels() -> (x_label, y_label)
    ssc_item_number : int, optional
        DFSU item number for SSC extraction.
    ssc_scale : {"log","normal"}, optional
        Color and y-axis scale for SSC.
    ssc_vmin, ssc_vmax : float | None, optional
        Colorbar min/max in mg/L. Defaults derive from `levels` or data.
    ssc_levels : list[float], optional
        Colorbar ticks in mg/L. Always includes the lower bound tick.
    ssc_cmap_name : str, optional
        Matplotlib colormap name.
    cbar_tick_decimals : int, optional
        Colorbar tick label precision.
    axis_tick_decimals : int, optional
        X/Y tick label precision for map and middle plot.
    pad_m : float, optional
        Padding in meters for the bbox around the transect.
    ssc_pixel_size_m : float, optional
        Rasterization resolution in meters.
    ssc_bottom_thresh : float, optional
        Values below this are transparent on the map. Also the lower bound for log plots.
    adcp_transect_color : str, optional
        Line color for the ADCP transect on the map.
    distance_bin_m : float, optional
        Bin width in meters for distance-binned bars.
    bar_width_scale : float, optional
        Relative bar width factor for middle plot (0–1).
    adcp_series_mode, adcp_series_target : str, optional
        ADCP aggregation configuration passed to `get_beam_series`.
    shapefile_layers : list[ShapefileLayer] | None, optional
        Optional overlays. Each layer must implement .plot(ax).

    Returns
    -------
    fig : matplotlib.figure.Figure
    axes : tuple[matplotlib.axes.Axes, matplotlib.axes.Axes, matplotlib.axes.Axes]
        (ax0, ax1, ax2)
    """
    # ============================= Internal constants (layout) =============================
    FIG_W, FIG_H = 6.5, 9.0
    LEFT, RIGHT = 0.06, 0.96
    TOP, BOTTOM = 0.98, 0.0
    HSPACE = 0.22
    CB_WIDTH = 0.012          # colorbar width (figure fraction)
    CB_GAP = 0.008            # gap between colorbar axis and figure edge

    # metadata spacing
    META_TOP_Y = 0.95
    META_SECTION_GAP = 0.20
    META_LINE_GAP = 0.16
    META_COL_X = [0.02, 0.43, 0.84]
    META_LEFT_OVERSHOOT = 0.10

    # ============================== Imports used inside ==============================
    import numpy as _np
    import pandas as _pd
    import matplotlib.pyplot as _plt
    from matplotlib import ticker as _mticker
    from matplotlib.colors import LogNorm as _LogNorm, Normalize as _Normalize
    from .plotting import PlottingShell as _PlottingShell

    # ============================== Data aliases ==============================
    xq = adcp.position.x
    yq = adcp.position.y
    t = adcp.time.ensemble_datetimes
    x_label, y_label = crs_helper.axis_labels()
    overlays = shapefile_layers or []

    # ============================== Figure shell ==============================
    fig, ax = _PlottingShell.subplots(
        figheight=FIG_H,
        figwidth=FIG_W,
        nrow=3,
        ncol=1,
        height_ratios=[1.00, 0.30, 0.22],
    )
    fig.subplots_adjust(left=LEFT, right=RIGHT, top=TOP, bottom=BOTTOM, hspace=HSPACE)
    ax0, ax1, ax2 = ax

    # ======================================================================
    # TOP — Model SSC raster (mg/L)
    # ======================================================================
    bbox = crs_helper.bbox_from_coords(xq, yq, pad_m=pad_m, from_crs=adcp.position.epsg)
    model_img = mt_model.rasterize_idw_bbox(
        item_number=ssc_item_number, bbox=bbox, t=t, pixel_size_m=ssc_pixel_size_m
    )
    model_ssc_img = _np.asarray(model_img[0], dtype=float) * 1000.0  # mg/L
    model_extent = model_img[1]

    finite = _np.isfinite(model_ssc_img)
    if not finite.any():
        raise ValueError("No finite SSC values in model raster.")

    auto_min = float(_np.nanmin(model_ssc_img[finite]))
    auto_max = float(_np.nanmax(model_ssc_img[finite]))
    cbar_min = ssc_vmin if ssc_vmin is not None else (min(ssc_levels) if ssc_levels else max(ssc_bottom_thresh, auto_min))
    cbar_max = ssc_vmax if ssc_vmax is not None else (max(ssc_levels) if ssc_levels else auto_max)
    cbar_min = max(ssc_bottom_thresh, cbar_min)
    if (not _np.isfinite(cbar_min)) or (not _np.isfinite(cbar_max)) or (cbar_min >= cbar_max):
        cbar_min, cbar_max = max(ssc_bottom_thresh, 0.01), max(ssc_bottom_thresh * 100.0, 100.0)

    cmap = _plt.get_cmap(ssc_cmap_name).copy()
    cmap.set_under(alpha=0.0)
    norm = _LogNorm(vmin=cbar_min, vmax=cbar_max, clip=True) if ssc_scale.lower() == "log" \
        else _Normalize(vmin=cbar_min, vmax=cbar_max, clip=True)

    im = ax0.imshow(
        model_ssc_img,
        extent=model_extent,
        origin="lower",
        cmap=cmap,
        norm=norm,
        interpolation="nearest",
    )

    for layer in overlays:
        layer.plot(ax0)

    ax0.plot(xq, yq, color=adcp_transect_color, lw=1.0, alpha=0.9)
    ax0.set_xlabel(x_label)
    ax0.set_ylabel(y_label)

    xmin, xmax, ymin, ymax = bbox
    ax0.set_xlim(xmin, xmax)
    ax0.set_ylim(ymin, ymax)
    ax0.set_aspect("equal", adjustable="box")
    ax0.set_title(f"Model Comparison to ADCP Transect {adcp.name}", fontsize=8)

    xy_fmt = _mticker.FuncFormatter(lambda v, pos: f"{v:.{axis_tick_decimals}f}")
    ax0.xaxis.set_major_formatter(xy_fmt)
    ax0.yaxis.set_major_formatter(xy_fmt)

    # --- Colorbar axis placed with a gap from the figure edge ---
    fig.canvas.draw()
    pos0 = ax0.get_position()
    ax0.set_position([pos0.x0, pos0.y0, pos0.width - (CB_WIDTH + CB_GAP), pos0.height])
    fig.canvas.draw()
    pos = ax0.get_position()

    cb_ax = fig.add_axes([pos.x1 + CB_GAP, pos.y0, CB_WIDTH, pos.height])
    cb = _plt.colorbar(im, cax=cb_ax)
    cb.ax.set_ylabel("Mean SSC During Transect (mg/L)", fontsize=7)
    cb.ax.tick_params(labelsize=6)

    def _fmt_num(val: float, nd: int) -> str:
        return f"{val:.{nd}f}"

    tick_list = sorted(set(list(ssc_levels) + [cbar_min])) if ssc_levels else [cbar_min]
    tick_list = [v for v in tick_list if cbar_min <= v <= cbar_max]
    if tick_list:
        cb.set_ticks(tick_list)
        cb.set_ticklabels([_fmt_num(v, cbar_tick_decimals) for v in tick_list])
    else:
        cb_ticks = _np.linspace(cbar_min, cbar_max, 5)
        cb.set_ticks(cb_ticks)
        cb.set_ticklabels([_fmt_num(v, cbar_tick_decimals) for v in cb_ticks])

    # Legend for ADCP track (solid white background, on top)
    from matplotlib.lines import Line2D
    
    _track = Line2D([0], [0], color=adcp_transect_color, lw=1.2)
    leg = ax0.legend([_track], ["ADCP track"],
                     loc="upper left", frameon=True, fontsize=7,
                     framealpha=1.0, fancybox=False)
    
    frame = leg.get_frame()
    frame.set_facecolor("white"); frame.set_alpha(1.0); frame.set_edgecolor("black")
    leg.set_zorder(1000); frame.set_zorder(1001)


    # ======================================================================
    # MIDDLE — Distance-binned SSC (Model vs ADCP)
    # ======================================================================
    model_transect = mt_model.extract_transect(xq, yq, t, item_number=ssc_item_number)
    model_ssc_ts = _np.asarray(model_transect[0], dtype=float) * 1000.0  # mg/L

    adcp_series = adcp.get_beam_series(
        field_name="suspended_solids_concentration",
        mode=adcp_series_mode,
        target=adcp_series_target,
    )
    adcp_ssc_ts = _np.asarray(adcp_series[0], dtype=float)

    dist_m = _np.asarray(adcp.position.distance, dtype=float).ravel()

    n = min(dist_m.size, model_ssc_ts.size, adcp_ssc_ts.size)
    dist_m = dist_m[:n]
    model_ssc_ts = model_ssc_ts[:n]
    adcp_ssc_ts = adcp_ssc_ts[:n]

    valid = _np.isfinite(dist_m) & _np.isfinite(model_ssc_ts) & _np.isfinite(adcp_ssc_ts)
    dist_m = dist_m[valid]
    model_ssc_ts = model_ssc_ts[valid]
    adcp_ssc_ts = adcp_ssc_ts[valid]
    if dist_m.size == 0:
        raise ValueError("No valid samples for distance binning.")

    dmax = float(_np.nanmax(dist_m))
    edges = _np.arange(0.0, dmax + distance_bin_m, distance_bin_m)
    if edges.size < 2:
        end_right = max(distance_bin_m, dmax if _np.isfinite(dmax) else distance_bin_m)
        edges = _np.array([0.0, end_right], dtype=float)

    centers = 0.5 * (edges[:-1] + edges[1:])
    nbins = edges.size - 1
    bin_idx = _np.clip(_np.digitize(dist_m, edges) - 1, 0, nbins - 1)

    def _binned_mean(values: _np.ndarray, idx: _np.ndarray, nb: int) -> _np.ndarray:
        sums = _np.bincount(idx, weights=values, minlength=nb).astype(float)
        cnts = _np.bincount(idx, minlength=nb).astype(float)
        out = _np.full(nb, _np.nan, dtype=float)
        ok = cnts > 0
        out[ok] = sums[ok] / cnts[ok]
        return out

    model_bin_mean = _binned_mean(model_ssc_ts, bin_idx, nbins)
    adcp_bin_mean = _binned_mean(adcp_ssc_ts, bin_idx, nbins)

    bar_w = max(0.5, distance_bin_m * bar_width_scale)

    if ssc_scale.lower() == "log":
        m_plot = _np.where(_np.isfinite(model_bin_mean) & (model_bin_mean > 0.0), model_bin_mean, _np.nan)
        a_plot = _np.where(_np.isfinite(adcp_bin_mean) & (adcp_bin_mean > 0.0), adcp_bin_mean, _np.nan)
        m_plot = _np.where(m_plot < cbar_min, cbar_min, m_plot)
        a_plot = _np.where(a_plot < cbar_min, cbar_min, a_plot)
    else:
        m_plot = model_bin_mean
        a_plot = adcp_bin_mean

    ax1.bar(
        centers - bar_w / 2.0,
        m_plot,
        width=bar_w,
        color=_PlottingShell.blue1,
        alpha=0.9,
        label="Model",
        align="center",
    )
    ax1.bar(
        centers + bar_w / 2.0,
        a_plot,
        width=bar_w,
        color=_PlottingShell.red1,
        alpha=0.9,
        label="ADCP",
        align="center",
    )

    ax1.set_xlim(edges[0], edges[-1])
    ax1.set_xlabel("Distance along transect (m)")
    ax1.set_ylabel("SSC (mg/L)")
    ax1.grid(alpha=0.3)
    ax1.legend(frameon=False, fontsize=7, ncol=2)
    
    # ticks exactly at each bin center
    ax1.set_xticks(centers)
    ax1.set_xticklabels([f"{c:.1f}" for c in centers])  # or use your desired precision
    ax1.set_xlim(edges[0], edges[-1])
    ax1.xaxis.set_major_formatter(_mticker.FuncFormatter(lambda v, pos: f"{v:.1f}"))
    #ax1.yaxis.set_major_formatter(_mticker.FuncFormatter(lambda v, pos: f"{v:.{tick_decimal_precision}f}"))


        
        
    # Match middle width to top; ensure rectangular, not square
    fig.canvas.draw()
    p0 = ax0.get_position()
    p1 = ax1.get_position()
    ax1.set_aspect("auto")
    ax1.set_anchor("W")
    try:
        ax1.set_box_aspect(None)
    except Exception:
        pass
    ax1.set_position([p0.x0, p1.y0, p0.width, p1.height])

    # Middle y-limit matches colorbar max for comparability
    if ssc_scale.lower() == "log":
        ax1.set_yscale("log")
        ax1.set_ylim(cbar_min, cbar_max)
    else:
        ax1.set_ylim(0.0, cbar_max)
    # --- Force plain numeric y-ticks on the middle subplot (log or linear) ---
    from matplotlib.ticker import FixedLocator, FuncFormatter
    
    # choose ticks from your colorbar 'levels' and clip to [cbar_min, cbar_max]
    _y_ticks = [v for v in ssc_levels if (v > 0) and (cbar_min <= v <= cbar_max)]
    if not _y_ticks:
        _y_ticks = np.geomspace(cbar_min, cbar_max, 5) if ssc_scale.lower() == "log" else np.linspace(cbar_min, cbar_max, 5)
    
    # apply AFTER ax1.set_yscale(...) and ax1.set_ylim(...)
    ax1.yaxis.set_major_locator(FixedLocator(_y_ticks))
    ax1.yaxis.set_minor_locator(FixedLocator([]))  # no minor ticks
    ax1.yaxis.set_major_formatter(FuncFormatter(lambda v, pos: f"{v:.{cbar_tick_decimals}f}"))
    
    
    # ======================================================================
    # BOTTOM — Metadata (three columns)
    # ======================================================================
    t_arr = _np.asarray(t)[:n]
    t_valid = t_arr[valid]
    t0 = _pd.to_datetime(t_valid.min())
    t1 = _pd.to_datetime(t_valid.max())
    dur_min = (t1 - t0).total_seconds() / 60.0

    mean_model = float(_np.nanmean(model_ssc_ts))
    mean_obs = float(_np.nanmean(adcp_ssc_ts))
    mean_err = float(_np.nanmean(model_ssc_ts - adcp_ssc_ts))

    ax2.clear()
    ax2.set_axis_off()

    def _fmt_time(dt):
        return dt.strftime("%d %b. %Y %H:%M")

    def _H(x, y, text):
        ax2.text(x, y, text, ha="left", va="top", fontsize=8, fontweight="bold", family="monospace")

    def _I(x, y, k, v):
        ax2.text(x, y, f"{k}: {v}", ha="left", va="top", fontsize=7, family="monospace")

    # Match bottom width to top, then overshoot further left
    fig.canvas.draw()
    p2 = ax2.get_position()
    ax2.set_aspect("auto")
    ax2.set_anchor("W")
    try:
        ax2.set_box_aspect(None)
    except Exception:
        pass
    new_x0 = p0.x0 - META_LEFT_OVERSHOOT
    new_w = min(p0.width + META_LEFT_OVERSHOOT, 1.0 - new_x0)
    ax2.set_position([new_x0, p2.y0, new_w, p2.height])

    y0_ax = META_TOP_Y
    sec = META_SECTION_GAP
    line = META_LINE_GAP
    cols_x = META_COL_X

    # Column 1
    x = cols_x[0]
    y = y0_ax
    _H(x, y, "Observation Aggregation")
    y -= sec
    _I(x, y, "Mode", adcp_series_mode)
    y -= line
    _I(x, y, "Target", adcp_series_target)

    # Column 2
    x = cols_x[1]
    y = y0_ax
    _H(x, y, "Transect Timing")
    y -= sec
    _I(x, y, "Start", _fmt_time(t0))
    y -= line
    _I(x, y, "End", _fmt_time(t1))
    y -= line
    _I(x, y, "Duration", f"{dur_min:.1f} min")

    # Column 3
    x = cols_x[2]
    y = y0_ax
    _H(x, y, "Model Fit")
    y -= sec
    _I(x, y, "Mean error", f"{mean_err:.{cbar_tick_decimals}f} mg/L")
    y -= line
    _I(x, y, "Mean observed", f"{mean_obs:.{cbar_tick_decimals}f} mg/L")
    y -= line
    _I(x, y, "Mean simulated", f"{mean_model:.{cbar_tick_decimals}f} mg/L")

    return fig, (ax0, ax1, ax2)

#%%

if __name__ == '__main__':
    
    # ============= Load Data ============
    # Project and configs
    xml_path = r'//usden1-stor.dhi.dk/Projects/61803553-05/Projects/Clean Project F3 2 Oct 2024.mtproj'
    project = XMLUtils(xml_path)
    adcp_cfgs = project.get_cfgs_from_survey(survey_name="20241002_F3(E)", survey_id=0, instrument_type="VesselMountedADCP")
    adcp_cfg = adcp_cfgs[9]

    # CRS + model path
    crs_helper = CRSHelper(project_crs=4326)
    model_fpath = r'//usden1-stor.dhi.dk/Projects/61803553-05/Models/F3/2024/10. October/MT/MTD20241002_1.dfsu'

    # Heavy objects
    mt_model = DfsuUtils2D(model_fpath, crs_helper=crs_helper)
    adcp     = ADCPDataset(adcp_cfg, name=adcp_cfg["name"])

    # ======================================
    # Example call
    shapefile_layers = [
        ShapefileLayer(
            path=r"\\usden1-stor.dhi.dk\Projects\61803553-05\GIS\SG Coastline\RD7550_CEx_SG_v20250509.shp",
            kind="polygon",
            crs_helper=crs_helper,
            poly_edgecolor="black",
            poly_linewidth=0.6,
            poly_facecolor="none",
            alpha=1.0,
            zorder=10,
        ),
        ShapefileLayer(
            path=r"\\usden1-stor.dhi.dk\Projects\61803553-05\GIS\F3\example point layer\points_labels.shp",
            kind="point",
            crs_helper=crs_helper,
            point_color="black",
            point_marker="o",
            point_markersize=10,
            alpha=1.0,
            zorder=14,
            label_text="Source Location",
            label_offset_pts= (1,1),
            label_fontsize=8,
            label_color="black",
            label_ha="left",
            label_va="center",
        )
    ]
    
    fig, (ax0, ax1, ax2) = mt_model_transect_comparison(
        mt_model=mt_model,
        adcp=adcp,
        crs_helper=crs_helper,
        ssc_item_number=1,
        ssc_scale="log",
        ssc_vmin=None,
        ssc_vmax=None,
        ssc_levels=[.01,0.1, 1, 10,50],
        ssc_cmap_name='jet',
        cbar_tick_decimals=2,
        axis_tick_decimals=3,
        pad_m=2500,
        ssc_pixel_size_m=10,
        ssc_bottom_thresh=0.001,
        adcp_transect_color=PlottingShell.red1,
        distance_bin_m=50,
        bar_width_scale=0.15,
        adcp_series_mode="range",
        adcp_series_target="mean",
        shapefile_layers=shapefile_layers,
    )
    
    plt.show()