"""Orchestration layer for a single SOLWEIG timestep.

The public entry point is :func:`calculate_core_fused`, which hands off
the full pipeline (shadows, ground temperature, GVF, thermal delay,
radiation, Tmrt) to a single fused Rust FFI call.

Pipeline::

    SVF resolution → Shadows → Ground temp → GVF
        → Thermal delay → Radiation → Tmrt
"""

from __future__ import annotations

from types import SimpleNamespace
from typing import TYPE_CHECKING

import numpy as np

_OUT_SHADOW = 1 << 0
_OUT_KDOWN = 1 << 1
_OUT_KUP = 1 << 2
_OUT_LDOWN = 1 << 3
_OUT_LUP = 1 << 4
_OUT_ALL = _OUT_SHADOW | _OUT_KDOWN | _OUT_KUP | _OUT_LDOWN | _OUT_LUP


def _arr_key(arr):
    if arr is None:
        return None
    return (arr.ctypes.data, arr.shape)


if TYPE_CHECKING:
    from .api import HumanParams, Location, PrecomputedData, SolweigResult, SurfaceData, ThermalState, Weather


def calculate_core_fused(
    surface: SurfaceData,
    location: Location,
    weather: Weather,
    human: HumanParams,
    precomputed: PrecomputedData | None,
    state: ThermalState | None,
    physics: SimpleNamespace | None,
    materials: SimpleNamespace | None,
    conifer: bool = False,
    wall_material: str | None = None,
    use_anisotropic_sky: bool = False,
    max_shadow_distance_m: float | None = None,
    return_state_copy: bool = True,
    requested_outputs: set[str] | None = None,
) -> SolweigResult:
    """
    Fused SOLWEIG calculation — single Rust FFI call per daytime timestep.

    Functionally identical to calculate_core() but orchestrates shadows, ground
    temperature, GVF, thermal delay, radiation, and Tmrt entirely within Rust,
    eliminating intermediate numpy allocations and FFI round-trips.

    This is the primary compute path used by ``calculate()``.
    Supports both isotropic and anisotropic (Perez) sky models.

    Args:
        surface: Surface/terrain data (DSM, vegetation, walls, land cover).
        location: Geographic location (latitude, longitude).
        weather: Weather conditions with derived sun position.
        human: Human parameters (height, posture, absorptivities).
        precomputed: Optional pre-computed data (SVF, shadow matrices).
        state: Optional thermal state for time-series (carries forward temperatures).
        physics: Optional physics parameters (vegetation transmissivity, etc.).
        materials: Optional material properties (albedo, emissivity by land cover).
        conifer: Treat vegetation as evergreen conifers (always leaf-on).
        wall_material: Wall material type ("brick", "concrete", "wood", "cobblestone").
        use_anisotropic_sky: Use anisotropic (Perez) diffuse sky model.
        max_shadow_distance_m: Maximum shadow reach in metres.
        return_state_copy: If True, return a deep-copied thermal state.
        requested_outputs: Set of output names to materialize (None = all).

    Returns:
        SolweigResult with Tmrt, shadow, radiation components, and updated state.
    """
    from .api import SolweigResult
    from .buffers import as_float32
    from .components.gvf import detect_building_mask
    from .components.shadows import compute_transmissivity
    from .components.svf_resolution import resolve_svf
    from .models.state import ThermalState
    from .physics.clearnessindex_2013b import clearnessindex_2013b
    from .physics.daylen import daylen
    from .physics.diffusefraction import diffusefraction
    from .rustalgos import pipeline

    # Ensure derived weather fields are computed (sun position, radiation split)
    if not weather._derived_computed:
        weather.compute_derived(location)

    # === Precompute (stays in Python) ===

    rows, cols = surface.dsm.shape
    pixel_size = surface.pixel_size

    # Valid pixel mask (True where all layers have finite data)
    # Computed once by SurfaceData.prepare(), or derived from DSM if missing
    valid_mask = surface.valid_mask
    valid_source = valid_mask if valid_mask is not None else surface.dsm
    valid_mask_key = _arr_key(valid_source)
    cache = surface._cache
    valid_mask_cache = cache.valid_mask_u8_cache
    if valid_mask_cache is not None and valid_mask_cache[0] == valid_mask_key:
        valid_mask_u8 = valid_mask_cache[1]
    else:
        if valid_mask is None:
            valid_mask = np.isfinite(surface.dsm)
        valid_mask_u8 = np.ascontiguousarray(valid_mask, dtype=np.uint8)
        cache.valid_mask_u8_cache = valid_mask_key, valid_mask_u8

    # Valid-bounds crop (ported from old main implementation):
    # trim heavy per-timestep compute to the minimal bounding rectangle of valid pixels.
    bbox_cache = cache.valid_bbox_cache
    if bbox_cache is not None and bbox_cache[0] == valid_mask_key:
        r0, r1, c0, c1 = bbox_cache[1]
    else:
        rows_any = np.any(valid_mask_u8 != 0, axis=1)
        cols_any = np.any(valid_mask_u8 != 0, axis=0)
        if not rows_any.any() or not cols_any.any():
            r0, r1, c0, c1 = 0, rows, 0, cols
        else:
            r_idx = np.flatnonzero(rows_any)
            c_idx = np.flatnonzero(cols_any)
            r0, r1 = int(r_idx[0]), int(r_idx[-1]) + 1
            c0, c1 = int(c_idx[0]), int(c_idx[-1]) + 1
        cache.valid_bbox_cache = valid_mask_key, (r0, r1, c0, c1)

    full_area = rows * cols
    crop_area = (r1 - r0) * (c1 - c0)
    use_crop = (r0, r1, c0, c1) != (0, rows, 0, cols) and crop_area < int(full_area * 0.98)
    crop_slice = (slice(r0, r1), slice(c0, c1))

    # Select which non-Tmrt outputs to materialize from Rust.
    if requested_outputs is None:
        output_mask = _OUT_ALL
    else:
        output_mask = 0
        if "shadow" in requested_outputs:
            output_mask |= _OUT_SHADOW
        if "kdown" in requested_outputs:
            output_mask |= _OUT_KDOWN
        if "kup" in requested_outputs:
            output_mask |= _OUT_KUP
        if "ldown" in requested_outputs:
            output_mask |= _OUT_LDOWN
        if "lup" in requested_outputs:
            output_mask |= _OUT_LUP

    # Land cover properties
    lc_props_key = (_arr_key(surface.land_cover), _arr_key(surface.albedo), _arr_key(surface.emissivity), id(materials))
    lc_props_cache = cache.land_cover_props_cache
    if lc_props_cache is not None and lc_props_cache[0] == lc_props_key:
        alb_grid, emis_grid, tgk_grid, tstart_grid, tmaxlst_grid = lc_props_cache[1]
    else:
        alb_grid, emis_grid, tgk_grid, tstart_grid, tmaxlst_grid = surface.get_land_cover_properties(materials)
        cache.land_cover_props_cache = lc_props_key, (alb_grid, emis_grid, tgk_grid, tstart_grid, tmaxlst_grid)

    # Vegetation inputs
    use_veg = surface.cdsm is not None
    cdsm = surface.cdsm if use_veg else None
    tdsm = surface.tdsm if use_veg else None
    if use_veg:
        pool = surface.get_buffer_pool()
        bush = pool.get_zeros("bush")
    else:
        bush = None

    # Wall inputs
    has_walls = surface.wall_height is not None and surface.wall_aspect is not None
    wall_ht = surface.wall_height if has_walls else None
    wall_asp = surface.wall_aspect if has_walls else None

    # Use full terrain relief for shadow ray termination so that mountain
    # ridges can correctly shadow valleys.  The horizontal reach is still
    # bounded by max_shadow_distance_m via max_index in Rust, so rays
    # don't run forever — they just won't terminate prematurely on the
    # vertical axis when terrain relief exceeds building heights.
    max_height = surface.max_height

    # SVF resolution (cached between timesteps)
    svf_bundle = resolve_svf(
        surface=surface,
        precomputed=precomputed,
    )

    # Vegetation transmissivity
    doy = weather.datetime.timetuple().tm_yday
    psi = compute_transmissivity(doy, physics, conifer)

    # Adjust svfbuveg for vegetation transmissivity (shortwave sees through canopy)
    # Without this, isotropic diffuse (drad), Kup, and Kdown treat vegetation as
    # fully opaque.  The anisotropic path already applies psi per sky patch via
    # diffsh(psi), and kside_veg applies psi per direction, but the scalar svfbuveg
    # used for isotropic diffuse and wall reflection was unadjusted.
    from .components.svf_resolution import adjust_svfbuveg_with_psi

    svf_bundle.svfbuveg = adjust_svfbuveg_with_psi(svf_bundle.svf, svf_bundle.svf_veg, psi, use_veg)

    # Wall material resolution
    tgk_wall = 0.37
    tstart_wall = -3.41
    tmaxlst_wall = 15.0
    albedo_wall = 0.20
    emis_wall = 0.90
    if wall_material is not None:
        from .loaders import resolve_wall_params

        tgk_wall, tstart_wall, tmaxlst_wall = resolve_wall_params(wall_material, materials)
    elif materials is not None:
        tgk_w = getattr(getattr(getattr(materials, "Ts_deg", None), "Value", None), "Walls", None)
        tstart_w = getattr(getattr(getattr(materials, "Tstart", None), "Value", None), "Walls", None)
        tmaxlst_w = getattr(getattr(getattr(materials, "TmaxLST", None), "Value", None), "Walls", None)
        if tgk_w is not None:
            tgk_wall = tgk_w
        if tstart_w is not None:
            tstart_wall = tstart_w
        if tmaxlst_w is not None:
            tmaxlst_wall = tmaxlst_w

    # Weather-derived scalars for ground temperature model
    _, _, _, snup = daylen(doy, location.latitude)
    dectime = (weather.datetime.hour + weather.datetime.minute / 60.0) / 24.0
    zen_deg = 90.0 - weather.sun_altitude

    # Clear-sky radiation for ground temperature CI correction
    zen_rad = zen_deg * (np.pi / 180.0)
    location_dict = {
        "latitude": location.latitude,
        "longitude": location.longitude,
        "altitude": 0.0,
    }
    i0, _, _, _, _ = clearnessindex_2013b(
        zen_rad,
        doy,
        weather.ta,
        weather.rh / 100.0,
        weather.global_rad,
        location_dict,
        -999.0,
    )
    if i0 > 0 and weather.sun_altitude > 0:
        rad_i0, rad_d0 = diffusefraction(i0, weather.sun_altitude, 1.0, weather.ta, weather.rh)
        rad_g0 = rad_i0 * np.sin(weather.sun_altitude * np.pi / 180.0) + rad_d0
    else:
        rad_g0 = 0.0

    # === Build Rust input structs ===

    ws = pipeline.WeatherScalars(
        sun_azimuth=float(weather.sun_azimuth),
        sun_altitude=float(weather.sun_altitude),
        sun_zenith=float(weather.sun_zenith),
        ta=float(weather.ta),
        rh=float(weather.rh),
        global_rad=float(weather.global_rad),
        direct_rad=float(weather.direct_rad),
        diffuse_rad=float(weather.diffuse_rad),
        altmax=float(weather.altmax),
        clearness_index=float(weather.clearness_index),
        dectime=float(dectime),
        snup=float(snup),
        rad_g0=float(rad_g0),
        zen_deg=float(zen_deg),
        psi=float(psi),
        is_daytime=weather.sun_altitude > 0,
        jday=int(weather.datetime.timetuple().tm_yday) if weather.datetime is not None else 180,
        patch_option=0,  # Set below if anisotropic
    )

    hs = pipeline.HumanScalars(
        height=float(human.height),
        abs_k=float(human.abs_k),
        abs_l=float(human.abs_l),
        is_standing=human.posture == "standing",
    )

    cs = pipeline.ConfigScalars(
        pixel_size=float(pixel_size),
        max_height=float(max_height),
        albedo_wall=float(albedo_wall),
        emis_wall=float(emis_wall),
        tgk_wall=float(tgk_wall),
        tstart_wall=float(tstart_wall),
        tmaxlst_wall=float(tmaxlst_wall),
        use_veg=use_veg,
        has_walls=has_walls,
        conifer=conifer,
        use_anisotropic=use_anisotropic_sky,
        max_shadow_distance_m=float(max_shadow_distance_m or 1000.0),
    )

    # Buildings mask for GVF (computed from DSM/land_cover/walls)
    buildings_key = (_arr_key(surface.dsm), _arr_key(surface.land_cover), _arr_key(wall_ht), float(pixel_size))
    buildings_cache = cache.buildings_mask_cache
    if buildings_cache is not None and buildings_cache[0] == buildings_key:
        buildings = buildings_cache[1]
    else:
        buildings = detect_building_mask(
            surface.dsm,
            surface.land_cover,
            wall_ht,
            pixel_size,
        )
        cache.buildings_mask_cache = buildings_key, buildings

    if surface.land_cover is not None:
        lc_grid_key = _arr_key(surface.land_cover)
        lc_grid_cache = cache.lc_grid_f32_cache
        if lc_grid_cache is not None and lc_grid_cache[0] == lc_grid_key:
            lc_grid = lc_grid_cache[1]
        else:
            lc_grid = surface.land_cover.astype(np.float32)
            cache.lc_grid_f32_cache = lc_grid_key, lc_grid
    else:
        lc_grid = None

    # GVF geometry cache: precompute on first daytime call, reuse on subsequent.
    # Keep separate caches for full-grid and cropped-grid execution.
    gvf_cache = None
    if has_walls:
        assert wall_asp is not None  # guaranteed by has_walls
        assert wall_ht is not None
        if use_crop:
            gvf_crop_key = (
                _arr_key(buildings),
                _arr_key(wall_asp),
                _arr_key(wall_ht),
                _arr_key(alb_grid),
                r0,
                r1,
                c0,
                c1,
                float(pixel_size),
                float(human.height),
                float(albedo_wall),
            )
            gvf_crop_cache = cache.gvf_geometry_cache_crop
            if gvf_crop_cache is not None and gvf_crop_cache[0] == gvf_crop_key:
                gvf_cache = gvf_crop_cache[1]
            else:
                gvf_cache = pipeline.precompute_gvf_cache(
                    as_float32(buildings[crop_slice]),
                    as_float32(wall_asp[crop_slice]),
                    as_float32(wall_ht[crop_slice]),
                    as_float32(alb_grid[crop_slice]),
                    float(pixel_size),
                    float(human.height),
                    float(albedo_wall),
                )
                cache.gvf_geometry_cache_crop = gvf_crop_key, gvf_cache
        else:
            gvf_cache = cache.gvf_geometry_cache
            if gvf_cache is None:
                gvf_cache = pipeline.precompute_gvf_cache(
                    as_float32(buildings),
                    as_float32(wall_asp),
                    as_float32(wall_ht),
                    as_float32(alb_grid),
                    float(pixel_size),
                    float(human.height),
                    float(albedo_wall),
                )
                cache.gvf_geometry_cache = gvf_cache

    # Anisotropic sky: Perez luminance, steradians, ASVF, and esky are now
    # computed inside the Rust pipeline (no Python round-trip). We only need
    # the shadow matrices and the patch_option.
    aniso_shmat = None
    aniso_vegshmat = None
    aniso_vbshmat = None

    if use_anisotropic_sky:
        shadow_mats = None
        if precomputed is not None and precomputed.shadow_matrices is not None:
            shadow_mats = precomputed.shadow_matrices
        elif surface.shadow_matrices is not None:
            shadow_mats = surface.shadow_matrices

        if shadow_mats is not None:
            ws.patch_option = shadow_mats.patch_option
            if use_crop:
                aniso_crop_key = (
                    _arr_key(shadow_mats._shmat_u8),
                    _arr_key(shadow_mats._vegshmat_u8),
                    _arr_key(shadow_mats._vbshmat_u8),
                    r0,
                    r1,
                    c0,
                    c1,
                )
                aniso_crop_cache = cache.aniso_shadow_crop_cache
                if aniso_crop_cache is not None and aniso_crop_cache[0] == aniso_crop_key:
                    aniso_shmat, aniso_vegshmat, aniso_vbshmat = aniso_crop_cache[1]
                else:
                    aniso_shmat = np.ascontiguousarray(shadow_mats._shmat_u8[crop_slice])
                    aniso_vegshmat = np.ascontiguousarray(shadow_mats._vegshmat_u8[crop_slice])
                    aniso_vbshmat = np.ascontiguousarray(shadow_mats._vbshmat_u8[crop_slice])
                    cache.aniso_shadow_crop_cache = aniso_crop_key, (aniso_shmat, aniso_vegshmat, aniso_vbshmat)
            else:
                # Keep original arrays to preserve stable pointers across timesteps.
                aniso_shmat = shadow_mats._shmat_u8
                aniso_vegshmat = shadow_mats._vegshmat_u8
                aniso_vbshmat = shadow_mats._vbshmat_u8

    # Thermal state (create initial if None)
    if state is None:
        state = ThermalState.initial((rows, cols))

    firstdaytime_int = int(state.firstdaytime)

    def _sel(arr):
        if arr is None:
            return None
        return arr[crop_slice] if use_crop else arr

    dsm_call = _sel(surface.dsm)
    cdsm_call = _sel(cdsm)
    tdsm_call = _sel(tdsm)
    bush_call = _sel(bush)
    wall_ht_call = _sel(wall_ht)
    wall_asp_call = _sel(wall_asp)
    svf_call = _sel(svf_bundle.svf)
    svf_n_call = _sel(svf_bundle.svf_directional.north)
    svf_e_call = _sel(svf_bundle.svf_directional.east)
    svf_s_call = _sel(svf_bundle.svf_directional.south)
    svf_w_call = _sel(svf_bundle.svf_directional.west)
    svf_veg_call = _sel(svf_bundle.svf_veg)
    svf_veg_n_call = _sel(svf_bundle.svf_veg_directional.north)
    svf_veg_e_call = _sel(svf_bundle.svf_veg_directional.east)
    svf_veg_s_call = _sel(svf_bundle.svf_veg_directional.south)
    svf_veg_w_call = _sel(svf_bundle.svf_veg_directional.west)
    svf_aveg_call = _sel(svf_bundle.svf_aveg)
    svf_aveg_n_call = _sel(svf_bundle.svf_aveg_directional.north)
    svf_aveg_e_call = _sel(svf_bundle.svf_aveg_directional.east)
    svf_aveg_s_call = _sel(svf_bundle.svf_aveg_directional.south)
    svf_aveg_w_call = _sel(svf_bundle.svf_aveg_directional.west)
    svfbuveg_call = _sel(svf_bundle.svfbuveg)
    svfalfa_call = _sel(svf_bundle.svfalfa)
    alb_call = _sel(alb_grid)
    emis_call = _sel(emis_grid)
    tgk_call = _sel(tgk_grid)
    tstart_call = _sel(tstart_grid)
    tmaxlst_call = _sel(tmaxlst_grid)
    buildings_call = _sel(buildings)
    lc_grid_call = _sel(lc_grid)
    valid_mask_call = _sel(valid_mask_u8)
    tgmap1_call = _sel(state.tgmap1)
    tgmap1_e_call = _sel(state.tgmap1_e)
    tgmap1_s_call = _sel(state.tgmap1_s)
    tgmap1_w_call = _sel(state.tgmap1_w)
    tgmap1_n_call = _sel(state.tgmap1_n)
    tgout1_call = _sel(state.tgout1)

    # === Call fused Rust pipeline ===

    result = pipeline.compute_timestep(
        # Scalar structs
        ws,
        hs,
        cs,
        # GVF geometry cache (None on first call triggers full GVF, then cached)
        gvf_cache,
        # Surface arrays
        as_float32(dsm_call),
        as_float32(cdsm_call) if cdsm_call is not None else None,
        as_float32(tdsm_call) if tdsm_call is not None else None,
        as_float32(bush_call) if bush_call is not None else None,
        as_float32(wall_ht_call) if wall_ht_call is not None else None,
        as_float32(wall_asp_call) if wall_asp_call is not None else None,
        # SVF arrays
        as_float32(svf_call),
        as_float32(svf_n_call),
        as_float32(svf_e_call),
        as_float32(svf_s_call),
        as_float32(svf_w_call),
        as_float32(svf_veg_call),
        as_float32(svf_veg_n_call),
        as_float32(svf_veg_e_call),
        as_float32(svf_veg_s_call),
        as_float32(svf_veg_w_call),
        as_float32(svf_aveg_call),
        as_float32(svf_aveg_n_call),
        as_float32(svf_aveg_e_call),
        as_float32(svf_aveg_s_call),
        as_float32(svf_aveg_w_call),
        as_float32(svfbuveg_call),
        as_float32(svfalfa_call),
        # Land cover property grids
        as_float32(alb_call),
        as_float32(emis_call),
        as_float32(tgk_call),
        as_float32(tstart_call),
        as_float32(tmaxlst_call),
        # Buildings mask + land cover
        as_float32(buildings_call),
        as_float32(lc_grid_call) if lc_grid_call is not None else None,
        # Anisotropic sky inputs (None for isotropic; Perez computed in Rust)
        aniso_shmat,
        aniso_vegshmat,
        aniso_vbshmat,
        # Thermal state
        firstdaytime_int,
        float(state.timeadd),
        float(state.timestep_dec),
        as_float32(tgmap1_call),
        as_float32(tgmap1_e_call),
        as_float32(tgmap1_s_call),
        as_float32(tgmap1_w_call),
        as_float32(tgmap1_n_call),
        as_float32(tgout1_call),
        # Valid pixel mask for early NaN exit
        valid_mask_call,
        output_mask,
    )

    # === Unpack result and update thermal state ===

    state.timeadd = result.timeadd
    if use_crop:
        state.tgmap1[crop_slice] = np.asarray(result.tgmap1)
        state.tgmap1_e[crop_slice] = np.asarray(result.tgmap1_e)
        state.tgmap1_s[crop_slice] = np.asarray(result.tgmap1_s)
        state.tgmap1_w[crop_slice] = np.asarray(result.tgmap1_w)
        state.tgmap1_n[crop_slice] = np.asarray(result.tgmap1_n)
        state.tgout1[crop_slice] = np.asarray(result.tgout1)
    else:
        state.tgmap1 = np.asarray(result.tgmap1)
        state.tgmap1_e = np.asarray(result.tgmap1_e)
        state.tgmap1_s = np.asarray(result.tgmap1_s)
        state.tgmap1_w = np.asarray(result.tgmap1_w)
        state.tgmap1_n = np.asarray(result.tgmap1_n)
        state.tgout1 = np.asarray(result.tgout1)

    if weather.is_daytime:
        state.firstdaytime = 0.0
    else:
        state.firstdaytime = 1.0
        state.timeadd = 0.0

    output_state = state.copy() if return_state_copy else state

    tmrt = np.asarray(result.tmrt)
    shadow = np.asarray(result.shadow) if result.shadow is not None else None
    kdown = np.asarray(result.kdown) if result.kdown is not None else None
    kup = np.asarray(result.kup) if result.kup is not None else None
    ldown = np.asarray(result.ldown) if result.ldown is not None else None
    lup = np.asarray(result.lup) if result.lup is not None else None

    if use_crop:

        def _uncrop(arr: np.ndarray | None) -> np.ndarray | None:
            if arr is None:
                return None
            full = np.full((rows, cols), np.nan, dtype=np.float32)
            full[crop_slice] = arr
            return full

        tmrt = _uncrop(tmrt)
        shadow = _uncrop(shadow)
        kdown = _uncrop(kdown)
        kup = _uncrop(kup)
        ldown = _uncrop(ldown)
        lup = _uncrop(lup)

    assert tmrt is not None  # tmrt is always computed
    return SolweigResult(
        tmrt=tmrt,
        shadow=shadow,
        kdown=kdown,
        kup=kup,
        ldown=ldown,
        lup=lup,
        utci=None,
        pet=None,
        state=output_state,
    )
