"""Time-series SOLWEIG calculation with thermal state management.

Provides :func:`_calculate_timeseries`, the internal implementation that
iterates over a list of :class:`~solweig.Weather` objects, carrying
thermal state (ground and wall temperatures) forward between timesteps.
Large rasters are transparently routed to the tiled processing path.

Users should call :func:`solweig._calculate_single` (the public entry point).
"""

from __future__ import annotations

import time
from collections.abc import Callable
from pathlib import Path
from types import SimpleNamespace
from typing import TYPE_CHECKING

import numpy as np

from .models import HumanParams, Location, ThermalState
from .output_async import AsyncGeoTiffWriter, async_output_enabled
from .postprocess import compute_utci_grid
from .progress import ProgressReporter
from .solweig_logging import get_logger
from .summary import GridAccumulator, TimeseriesSummary

logger = get_logger(__name__)


def _precompute_weather(weather_series: list, location: Location) -> None:
    """
    Pre-compute derived weather values for all timesteps efficiently.

    Optimizations:
    1. Compute max sun altitude (altmax) only once per unique day
    2. Pre-assign altmax to Weather objects to skip the 96-iteration loop

    This reduces compute_derived() from O(96) iterations to O(1) per timestep
    when multiple timesteps share the same day.

    Args:
        weather_series: List of Weather objects to process
        location: Geographic location for sun position calculations
    """
    if not weather_series:
        return

    from datetime import timedelta

    from .physics import sun_position as sp

    location_dict = location.to_sun_position_dict()

    # Step 1: Compute altmax once per unique day
    altmax_cache = {}  # date -> altmax

    for weather in weather_series:
        day = weather.datetime.date()
        if day not in altmax_cache:
            # Compute max sun altitude for this day (iterate in 15-min intervals)
            ymd = weather.datetime.replace(hour=0, minute=0, second=0, microsecond=0)
            sunmaximum = -90.0
            fifteen_min = 15.0 / 1440.0  # 15 minutes as fraction of day

            for step in range(96):  # 24 hours * 4 (15-min intervals)
                step_time = ymd + timedelta(days=step * fifteen_min)
                time_dict_step = {
                    "year": step_time.year,
                    "month": step_time.month,
                    "day": step_time.day,
                    "hour": step_time.hour,
                    "min": step_time.minute,
                    "sec": 0,
                    "UTC": location.utc_offset,
                }
                sun_step = sp.sun_position(time_dict_step, location_dict)
                zenith_step = sun_step["zenith"]
                zenith_val = (
                    float(np.asarray(zenith_step).flat[0]) if hasattr(zenith_step, "__iter__") else float(zenith_step)
                )
                altitude_step = 90.0 - zenith_val
                if altitude_step > sunmaximum:
                    sunmaximum = altitude_step

            altmax_cache[day] = max(sunmaximum, 0.0)

    # Step 2: Pre-assign altmax to each weather object
    for weather in weather_series:
        day = weather.datetime.date()
        weather.precomputed_altmax = altmax_cache[day]

    # Step 3: Compute derived values (now fast since altmax is cached)
    for weather in weather_series:
        if not weather._derived_computed:
            weather.compute_derived(location)

    # Step 4: Carry forward clearness index at night.
    # UMEP Python does not recompute CI when the sun is below the horizon;
    # instead the last daytime CI value persists into the nighttime hours.
    # This matters for the Ldown cloud-correction: if the last daytime CI
    # was < 0.95, the correction raises nighttime Ldown (more cloud → more
    # downwelling longwave).  Without carry-forward, CI defaults to 1.0
    # at night and the correction is never triggered.
    last_daytime_ci = 1.0
    for weather in weather_series:
        if weather.sun_altitude > 0 and weather.global_rad > 0:
            last_daytime_ci = weather.clearness_index
        else:
            weather.clearness_index = last_daytime_ci


if TYPE_CHECKING:
    from .models import (
        ModelConfig,
        PrecomputedData,
        SurfaceData,
        Weather,
    )


def _calculate_timeseries(
    surface: SurfaceData,
    weather_series: list[Weather],
    location: Location | None = None,
    config: ModelConfig | None = None,
    human: HumanParams | None = None,
    precomputed: PrecomputedData | None = None,
    use_anisotropic_sky: bool | None = None,
    conifer: bool = False,
    physics: SimpleNamespace | None = None,
    materials: SimpleNamespace | None = None,
    wall_material: str | None = None,
    max_shadow_distance_m: float | None = None,
    tile_workers: int | None = None,
    tile_queue_depth: int | None = None,
    prefetch_tiles: bool | None = None,
    *,
    output_dir: str | Path,
    outputs: list[str] | None = None,
    heat_thresholds_day: list[float] | None = None,
    heat_thresholds_night: list[float] | None = None,
    progress_callback: Callable[[int, int], None] | None = None,
) -> TimeseriesSummary:
    """
    Calculate Tmrt for a time series of weather data.

    Returns a :class:`TimeseriesSummary` with aggregated per-pixel grids
    (mean/max/min Tmrt and UTCI, sun/shade hours, heat-stress exceedance).
    Per-timestep arrays are written to disk and freed after each step.

    Maintains thermal state across timesteps for accurate surface temperature
    modeling with thermal inertia (TsWaveDelay_2015a).

    Large rasters are automatically processed using overlapping tiles to
    manage memory. The tile size and buffer distance are computed dynamically
    from available GPU/RAM resources and the maximum building height in the DSM,
    ensuring accurate shadows at tile boundaries without wasting overlap on
    short buildings.

    This is a convenience function that manages state automatically. For custom
    control over state, use calculate() directly with the state parameter.

    Args:
        surface: Surface/terrain data (DSM required, CDSM/DEM optional).
        weather_series: List of Weather objects in chronological order.
            The datetime of each Weather object determines the timestep size.
        location: Geographic location (lat, lon, UTC offset). If None, automatically
            extracted from surface's CRS metadata.
        config: Model configuration object providing base settings.
            Explicit parameters override config values when provided.
        human: Human body parameters (absorption, posture, weight, height, etc.).
            If None, uses config.human or HumanParams defaults.
        precomputed: Pre-computed SVF and/or shadow matrices. Optional.
        use_anisotropic_sky: Use anisotropic sky model.
            If None, uses config.use_anisotropic_sky (default True).
        conifer: Treat vegetation as evergreen conifers (always leaf-on). Default False.
        physics: Physics parameters (Tree_settings, Posture geometry) from load_physics().
            Site-independent scientific constants. If None, uses config.physics or bundled defaults.
        materials: Material properties (albedo, emissivity per landcover class) from load_materials().
            Site-specific landcover parameters. Only needed if surface has land_cover grid.
            If None, uses config.materials.
        wall_material: Wall material type for temperature model.
            One of "brick", "concrete", "wood", "cobblestone" (case-insensitive).
            If None (default), uses generic wall params from materials JSON.
        max_shadow_distance_m: Maximum shadow reach in metres (default 1000.0).
            Caps horizontal shadow ray distance and serves as the tile overlap
            buffer for automatic tiled processing of large rasters. If None,
            uses config.max_shadow_distance_m or 1000.0.
        tile_workers: Number of workers for tiled orchestration. If None, uses
            config.tile_workers or adaptive default.
        tile_queue_depth: Extra queued tile tasks beyond active workers. If None,
            uses config.tile_queue_depth or a runtime default.
        prefetch_tiles: Whether to prefetch extra tile tasks beyond active workers.
            If None, uses config.prefetch_tiles or runtime auto-selection.
        output_dir: Working directory for all output. Summary grids are always
            saved to ``output_dir/summary/``. Per-timestep GeoTIFFs are saved
            when ``outputs`` is specified.
        outputs: Which per-timestep outputs to save as GeoTIFFs (e.g., ["tmrt", "shadow"]).
            If None, only summary grids are saved.
        heat_thresholds_day: UTCI thresholds (°C) for daytime exceedance hours.
            Default ``[32, 38]`` (strong / very strong heat stress).
        heat_thresholds_night: UTCI thresholds (°C) for nighttime exceedance hours.
            Default ``[26]`` (tropical night threshold).
        progress_callback: Optional callback(current_step, total_steps) called after
            each timestep. If None, a tqdm progress bar is shown automatically.

    Returns:
        :class:`TimeseriesSummary` with aggregated grids and metadata.

    Example:
        # Summary grids only
        summary = _calculate_timeseries(
            surface=surface,
            weather_series=weather_list,
            output_dir="output/",
        )
        print(summary.tmrt_mean, summary.utci_hours_above[32])

        # With per-timestep GeoTIFFs
        summary = _calculate_timeseries(
            surface=surface,
            weather_series=weather_list,
            output_dir="output/",
            outputs=["tmrt", "shadow"],
        )
    """
    if not weather_series:
        return TimeseriesSummary.empty()

    anisotropic_requested_explicitly = use_anisotropic_sky is True

    # Auto-extract location from surface if not provided
    if location is None:
        logger.warning(
            "Location not provided - auto-extracting from surface CRS.\n"
            "⚠️  UTC offset will default to 0 if not specified, which may cause incorrect sun positions.\n"
            "   Recommend: provide location explicitly with correct UTC offset."
        )
        location = Location.from_surface(surface)

    # Build effective configuration: explicit params override config
    from ._orchestration import resolve_config_params

    _resolved = resolve_config_params(
        config=config,
        use_anisotropic_sky=use_anisotropic_sky,
        human=human,
        physics=physics,
        materials=materials,
        outputs=outputs,
        max_shadow_distance_m=max_shadow_distance_m,
        tile_workers=tile_workers,
        tile_queue_depth=tile_queue_depth,
        prefetch_tiles=prefetch_tiles,
    )
    use_anisotropic_sky = _resolved["use_anisotropic_sky"]
    human = _resolved["human"]
    physics = _resolved["physics"]
    materials = _resolved["materials"]
    outputs = _resolved["outputs"]
    effective_max_shadow = _resolved["max_shadow_distance_m"]
    tile_workers = _resolved["tile_workers"]
    tile_queue_depth = _resolved["tile_queue_depth"]
    prefetch_tiles = _resolved["prefetch_tiles"]
    anisotropic_arg = (
        use_anisotropic_sky if (anisotropic_requested_explicitly or use_anisotropic_sky is False) else None
    )

    # Fill NaN in surface layers (idempotent — skipped if already done)
    surface.fill_nan()

    # Auto-tile large rasters transparently
    from .tiling import _should_use_tiling

    if _should_use_tiling(surface.shape[0], surface.shape[1]):
        from .tiling import _calculate_timeseries_tiled

        logger.info(
            f"Raster size {surface.dsm.shape[1]}×{surface.dsm.shape[0]} exceeds tiling threshold — "
            "switching to tiled processing."
        )
        return _calculate_timeseries_tiled(
            surface=surface,
            weather_series=weather_series,
            location=location,
            human=human,
            precomputed=precomputed,
            use_anisotropic_sky=anisotropic_arg,
            conifer=conifer,
            physics=physics,
            materials=materials,
            wall_material=wall_material,
            max_shadow_distance_m=effective_max_shadow,
            tile_workers=tile_workers,
            tile_queue_depth=tile_queue_depth,
            prefetch_tiles=prefetch_tiles,
            output_dir=output_dir,
            outputs=outputs,
            heat_thresholds_day=heat_thresholds_day,
            heat_thresholds_night=heat_thresholds_night,
            progress_callback=progress_callback,
        )

    # Log configuration summary
    logger.info("=" * 60)
    logger.info("Starting SOLWEIG timeseries calculation")
    logger.info(f"  Grid size: {surface.dsm.shape[1]}×{surface.dsm.shape[0]} pixels")
    logger.info(f"  Timesteps: {len(weather_series)}")
    start_str = weather_series[0].datetime.strftime("%Y-%m-%d %H:%M")
    end_str = weather_series[-1].datetime.strftime("%Y-%m-%d %H:%M")
    logger.info(f"  Period: {start_str} → {end_str}")
    logger.info(f"  Location: {location.latitude:.2f}°N, {location.longitude:.2f}°E")

    options = []
    if use_anisotropic_sky:
        options.append("anisotropic sky")
    if precomputed is not None:
        options.append("precomputed SVF")
    if options:
        logger.info(f"  Options: {', '.join(options)}")

    if outputs:
        logger.info(f"  Auto-save: {output_dir} ({', '.join(outputs)})")
    else:
        logger.info(f"  Output dir: {output_dir} (summary only)")
    logger.info("=" * 60)

    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    # Determine which arrays the Rust compute needs to return.
    # tmrt + shadow are always needed for the summary accumulator.
    requested_outputs: set[str] = {"tmrt", "shadow"}
    if outputs:
        requested_outputs |= set(outputs)

    from .api import _calculate_single

    # Pre-compute derived weather values in parallel (sun position, radiation split)
    # This is ~4x faster than computing sequentially in the main loop
    logger.info("Pre-computing sun positions and radiation splits...")
    precompute_start = time.time()
    _precompute_weather(weather_series, location)
    precompute_time = time.time() - precompute_start
    logger.info(f"  Pre-computed {len(weather_series)} timesteps in {precompute_time:.1f}s")

    processed_steps = 0
    state = ThermalState.initial(surface.shape)

    # Pre-calculate timestep size from first two entries (matching runner behavior)
    # The runner uses a fixed timestep_dec for all iterations, calculated upfront
    if len(weather_series) >= 2:
        dt0 = weather_series[0].datetime
        dt1 = weather_series[1].datetime
        state.timestep_dec = (dt1 - dt0).total_seconds() / 86400.0
        _timestep_hours = (dt1 - dt0).total_seconds() / 3600.0
    else:
        _timestep_hours = 1.0

    # Grid accumulator for summary statistics
    _accumulator = GridAccumulator(
        shape=surface.shape,
        heat_thresholds_day=heat_thresholds_day if heat_thresholds_day is not None else [32.0, 38.0],
        heat_thresholds_night=heat_thresholds_night if heat_thresholds_night is not None else [26.0],
        timestep_hours=_timestep_hours,
    )

    # Pre-create buffer pool for array reuse across timesteps
    _ = surface.get_buffer_pool()

    # Set up progress reporting (caller callback suppresses tqdm)
    n_steps = len(weather_series)
    _progress = None if progress_callback is not None else ProgressReporter(total=n_steps, desc="SOLWEIG timeseries")
    use_async_output = bool(outputs) and async_output_enabled()
    _writer = AsyncGeoTiffWriter(output_dir=output_path, surface=surface) if use_async_output else None

    # Start timing
    start_time = time.time()

    try:
        for i, weather in enumerate(weather_series):
            # Process timestep
            result = _calculate_single(
                surface=surface,
                location=location,
                weather=weather,
                human=human,
                precomputed=precomputed,
                use_anisotropic_sky=anisotropic_arg,
                conifer=conifer,
                state=state,
                physics=physics,
                materials=materials,
                wall_material=wall_material,
                max_shadow_distance_m=effective_max_shadow,
                return_state_copy=False,
                _requested_outputs=requested_outputs,
            )

            # Carry forward state to next timestep
            if result.state is not None:
                state = result.state
                result.state = None  # Free state arrays (~23 MB); state managed externally

            # Update grid accumulator (before potential array release)
            _accumulator.update(result, weather, compute_utci_fn=compute_utci_grid)

            from ._orchestration import process_timestep_result

            process_timestep_result(
                result,
                weather,
                outputs,
                human,
                _writer,
                output_dir,
                surface=surface,
            )
            processed_steps += 1

            # Report progress
            if progress_callback is not None:
                progress_callback(i + 1, n_steps)
            elif _progress is not None:
                _progress.update(1)
    finally:
        if _progress is not None:
            _progress.close()
        if _writer is not None:
            _writer.close()

    # Finalize summary
    from ._orchestration import finalize_summary

    return finalize_summary(
        _accumulator,
        surface,
        processed_steps=processed_steps,
        start_time=start_time,
        location=location,
        weather_series=weather_series,
        human=human,
        physics=physics,
        materials=materials,
        use_anisotropic_sky=use_anisotropic_sky,
        conifer=conifer,
        output_dir=output_dir,
        outputs=outputs,
    )
