import numpy as np
from dataclasses import dataclass, field
import warnings
from .units import convert_units, LengthUnits
from .geometry.occlusion import compute_transmission

np.seterr(divide="ignore", invalid="ignore")


@dataclass(frozen=True)
class LampCacheEntry:
    base_values: np.ndarray | None
    values: np.ndarray | None
    calc_state: tuple | None
    update_state: tuple | None


@dataclass(frozen=True)
class SurfaceCacheEntry:
    form_factors: np.ndarray | None
    theta_zone: np.ndarray | None
    values: np.ndarray | None
    surface_calc_state: tuple | None
    surface_update_state: tuple | None


@dataclass(frozen=True)
class ZoneCache:
    lamp_cache: dict = field(default_factory=dict)
    surface_cache: dict = field(default_factory=dict)
    calc_state: tuple | None = None
    update_state: tuple | None = None
    scene_geometry_state: tuple | None = None

    def needs_recalc(self, calc_state, lamp_id, lamp_calc_state):
        if calc_state != self.calc_state:
            return True  # zone has changed
        if self.lamp_cache.get(lamp_id) is None:
            return True  # new lamp
        if self.lamp_cache.get(lamp_id).calc_state != lamp_calc_state:
            return True  # lamp has changed
        return False

    def needs_update(self, update_state, lamp_id, lamp_update_state):
        if update_state != self.update_state:
            return True  # zone has changed
        if self.lamp_cache.get(lamp_id) is None:
            return True  # new lamp
        if self.lamp_cache.get(lamp_id).update_state != lamp_update_state:
            return True
        return False

    def needs_surface_recalc(self, surface_id, surface_calc_state, zone_calc_state,
                             scene_geometry_state):
        """Check if form factors need recomputing for a surface→zone pair."""
        if zone_calc_state != self.calc_state:
            return True
        if scene_geometry_state != self.scene_geometry_state:
            return True  # some surface moved → occlusion changed
        entry = self.surface_cache.get(surface_id)
        if entry is None:
            return True
        if entry.surface_calc_state != surface_calc_state:
            return True
        return False

    def needs_surface_update(self, surface_id, surface_update_state):
        """Check if reflected values need recomputing (surface incidence changed)."""
        entry = self.surface_cache.get(surface_id)
        if entry is None:
            return True
        if entry.surface_update_state != surface_update_state:
            return True
        return False

    def base_values(self, lamp_id):
        """convenience attribute for fetching lamp base values"""
        entry = self.lamp_cache.get(lamp_id)
        if entry is not None:
            return entry.base_values
        return None

    def values(self, lamp_id):
        """convenience attribute for fetching lamp values"""
        entry = self.lamp_cache.get(lamp_id)
        if entry is not None:
            return entry.values
        return None

    def by_wavelength(self, lamps, reduce=None):
        """Return {wavelength: summed values} for cached lamps.

        If reduce is given (e.g. np.mean, np.max), apply it to each
        wavelength's summed array before returning.
        """
        result = {}
        for key, lamp in lamps.items():
            entry = self.lamp_cache.get(key)
            if entry is None or entry.values is None:
                continue
            if lamp.wavelength is None:
                msg = (
                    f"{lamp.name} ({key}) has an undefined wavelength. "
                    f"Its contribution will not be counted."
                )
                warnings.warn(msg, stacklevel=2)
                continue
            wv = lamp.wavelength
            if wv in result:
                result[wv] = result[wv] + entry.values
            else:
                result[wv] = entry.values.copy()
        if reduce is not None:
            result = {wv: reduce(vals) for wv, vals in result.items()}
        return result


class LightingCalculator:
    """
    Performs all computations for a calculation zone: direct irradiance,
    occlusion, and reflected contributions.
    """

    def __init__(self):
        self.cache = ZoneCache()

    def compute(self, lamps, zv, surfaces=None, enable_occlusion=True, hard=False):
        """
        Calculate and return direct irradiance values at all coordinate points.
        """

        if len(lamps) == 0:
            self.cache = ZoneCache()
            return np.zeros(zv.num_points, dtype="float32")

        occluders = surfaces if enable_occlusion else None

        lamp_cache = {}
        for lamp_id, lamp in lamps.items():
            # potentially expensive
            base_values = self.calculate_lamp(lamp, zv, surfaces=occluders, hard=hard)
            # always cheap
            values = self.apply_filters(lamp, base_values.copy(), zv)
            lamp_cache[lamp_id] = LampCacheEntry(
                base_values=base_values,
                values=values,
                calc_state=lamp.calc_state,
                update_state=lamp.update_state,
            )

        # build scene geometry state for occlusion cache invalidation
        scene_geo = None
        if surfaces:
            scene_geo = tuple(s.calc_state for s in surfaces.values())

        # update cache
        self.cache = ZoneCache(
            lamp_cache=lamp_cache,
            surface_cache=self.cache.surface_cache,  # preserve across compute calls
            calc_state=zv.calc_state,
            update_state=zv.update_state,
            scene_geometry_state=scene_geo,
        )

        # sum across lamp values
        return self.aggregate(lamps, zv)

    def compute_reflectance(self, surfaces, zv, enable_occlusion=True, hard=False):
        """Compute total reflected contribution from all surfaces to a zone."""
        if not surfaces:
            return np.zeros(zv.num_points, dtype="float32")

        scene_geo = self.cache.scene_geometry_state

        surface_cache = {}
        total = np.zeros(zv.num_points, dtype="float32")

        for surface_id, surface in surfaces.items():
            if surface.R == 0 or surface.plane.values is None:
                continue

            RECALC = self.cache.needs_surface_recalc(
                surface_id, surface.calc_state, zv.calc_state, scene_geo
            ) or hard
            UPDATE = self.cache.needs_surface_update(
                surface_id, surface.update_state
            ) or RECALC

            if RECALC:
                form_factors, theta_zone = surface._calculate_coordinates(zv)
                if enable_occlusion:
                    transmission = compute_transmission(
                        surface.plane.coords, zv.coords, surfaces, exclude=surface_id
                    )
                    form_factors = form_factors * transmission.reshape(form_factors.shape)
            else:
                entry = self.cache.surface_cache[surface_id]
                form_factors = entry.form_factors
                theta_zone = entry.theta_zone

            if UPDATE:
                values = surface._calculate_values(form_factors, theta_zone, zv)
            else:
                values = self.cache.surface_cache[surface_id].values

            surface_cache[surface_id] = SurfaceCacheEntry(
                form_factors=form_factors,
                theta_zone=theta_zone,
                values=values,
                surface_calc_state=surface.calc_state,
                surface_update_state=surface.update_state,
            )

            total += (values * surface.R).reshape(*zv.num_points).astype("float32")

        # update cache with new surface entries
        self.cache = ZoneCache(
            lamp_cache=self.cache.lamp_cache,
            surface_cache=surface_cache,
            calc_state=self.cache.calc_state,
            update_state=self.cache.update_state,
            scene_geometry_state=scene_geo,
        )

        return total.astype("float32")

    def _to_meters(self, R, lamp):
        """Convert distance array to meters for inverse-square calculation."""
        if lamp.surface.units != LengthUnits.METERS:
            return np.array(convert_units(lamp.surface.units, "meters", *R))
        return R

    def _irradiance_at(self, lamp, coords):
        """Compute raw irradiance at arbitrary coordinates for a single lamp."""
        rel_coords = coords - lamp.surface.position
        Theta, Phi, R = lamp.transform_to_lamp(rel_coords, which="polar")
        R = self._to_meters(R, lamp)
        phot = lamp.ies.photometry.interpolated()
        values = phot.get_intensity(Theta, Phi) / R ** 2

        if lamp.nearfield:
            phot_dist_m = float(self._to_meters(
                np.array([lamp.surface.photometric_distance]), lamp
            ).ravel()[0])
            near_idx = np.where(R < phot_dist_m)
            if len(near_idx[0]) > 0:
                values[near_idx] = 0
                points = lamp.surface.surface_points
                imap = lamp.surface.intensity_map.reshape(-1)
                n_src = len(points)
                for point, mult in zip(points, imap):
                    rel = coords - point
                    Th, Ph, Rn = lamp.transform_to_lamp(rel, which="polar")
                    Rn_m = self._to_meters(Rn[near_idx], lamp)
                    nv = phot.get_intensity(Th[near_idx], Ph[near_idx]) / Rn_m ** 2
                    values[near_idx] += nv * mult / n_src

        return values, R

    def _cell_average_near_lamp(self, lamp, values, zv, R):
        """
        Replace near-lamp point-samples with median of cell sub-samples.
        NOTE: This is an improvement over Visual, which does not do any management of
        source-to-course-point blowups; which makes sense since Visual is rarely used
        with the average room fluence rate volume, where this fix actually matters.
        This note is here to clarify that because of this patch, guv-calcs outputs for
        fluence rate will often be lower compared to Visual--however, we currently believe
        that the guv-calcs solution is correct.
        """
        spacings = []
        for axis in range(3):
            u = np.unique(zv.coords[:, axis])
            if len(u) > 1:
                spacings.append(np.min(np.diff(u)))
        if not spacings:
            return values
        min_spacing = min(spacings)
        min_spacing_m = float(self._to_meters(
            np.array([min_spacing]), lamp
        ).ravel()[0])

        close_idx = np.where(R < min_spacing_m)[0]
        if len(close_idx) == 0:
            return values

        # Generate sub-sample offsets within each cell (5 per active dimension)
        n_sub = 5
        offset_arrays = []
        for axis in range(3):
            u = np.unique(zv.coords[:, axis])
            if len(u) > 1:
                s = np.min(np.diff(u))
                offset_arrays.append(
                    np.linspace(-s / 2, s / 2, n_sub, endpoint=False) + s / (2 * n_sub)
                )
            else:
                offset_arrays.append(np.array([0.0]))

        grids = np.meshgrid(*offset_arrays, indexing='ij')
        sub_offsets = np.column_stack([g.ravel() for g in grids])

        for idx in close_idx:
            sub_coords = zv.coords[idx] + sub_offsets
            sub_vals, _ = self._irradiance_at(lamp, sub_coords)
            valid = np.isfinite(sub_vals)
            if np.any(valid):
                values[idx] = np.median(sub_vals[valid])

        return values

    def calculate_lamp(self, lamp, zv, surfaces=None, hard=False):
        """Calculate the zone values for a single lamp."""

        RECALC = self.cache.needs_recalc(
            calc_state=zv.calc_state,
            lamp_id=lamp.lamp_id,
            lamp_calc_state=lamp.calc_state,
        )
        if hard or RECALC:
            values, R = self._irradiance_at(lamp, zv.coords)

            # Cell-average near-lamp points to prevent grid-alignment spikes
            values = self._cell_average_near_lamp(lamp, values, zv, R)

            # apply occlusion from room surfaces
            # use all surface points as sources for soft (penumbral) shadows
            if surfaces:
                source_pts = lamp.surface.surface_points
                if source_pts.ndim == 1:
                    source_pts = source_pts[None, :]
                transmission = compute_transmission(
                    source_pts, zv.coords, surfaces
                )
                avg_transmission = transmission.mean(axis=0)
                values *= avg_transmission.reshape(values.shape)

            if any(~np.isfinite(values)):  # mask any nans and infs near light source
                values = np.ma.masked_invalid(values)

            return values.astype("float32")
        # if no recalculation required use cached version
        return self.cache.base_values(lamp.lamp_id)

    def apply_filters(self, lamp, values, zv):
        """
        update the values of a single lamp based on the calc zone properties,
        but which don't require a full recalculation
        """

        if not zv.is_volume():
            if zv.has_view_mode():
                theta, elevation, azimuth = _compute_view_angles(lamp, zv)
                values = apply_plane_filters(values, theta, zv, elevation, azimuth)
            else:
                rel_coords = zv.coords - lamp.surface.position
                x, y, z = (rel_coords @ zv.basis).T
                r = np.sqrt(x ** 2 + y ** 2 + z ** 2)
                theta = np.arccos(-z / r)
                values = apply_plane_filters(values, theta, zv)

        # TODO: This should maybe actually be in Lamp/Photometry
        values = values * lamp.intensity_units.factor

        return values.reshape(*zv.num_points)

    def aggregate(self, lamps, zv):
        """sum across lamp_values"""

        lamp_values = [self.cache.values(lamp_id) for lamp_id in lamps.keys()]

        if not zv.is_volume() and zv.fov_horiz < 360 and len(lamps) > 1 and not zv.has_view_mode():
            base_values = self.calculate_horizontal_fov(lamps, lamp_values, zv)
        else:
            base_values = sum(lamp_values)
        return base_values.reshape(*zv.num_points).astype("float32")

    def calculate_horizontal_fov(self, lamps, lamp_values, zv):
        """
        Vectorized function to compute the largest possible value for all lamps
        within a horizontal view field.
        """

        # Compute relative coordinates: Shape (num_points, num_lamps, 3)
        lamp_positions = np.array([lamp.surface.position for lamp in lamps.values()])
        rel_coords = zv.coords[:, None, :] - lamp_positions[None, :, :]
        rel_coords = rel_coords @ zv.basis

        # Calculate horizontal angles (in degrees)
        angles = np.degrees(np.arctan2(rel_coords[..., 1], rel_coords[..., 0]))
        angles = angles % 360  # Wrap; Shape (N, M)
        angles = angles[:, :, None]  # Expand; Shape (N, M, 1)

        # Compute pairwise angular differences for all rows
        diffs = np.abs(angles - angles.transpose(0, 2, 1))
        diffs = np.minimum(diffs, 360 - diffs)  # Wrap angular differences to [0, 180]

        # Create the adjacency mask for each pair within 180 degrees
        adjacency = diffs <= zv.fov_horiz / 2  # Shape (N, M, M)

        # current values to be transformed
        values = np.array([val.reshape(-1) for val in lamp_values]).T
        # Sum the values for all connected components (using the adjacency mask)
        value_sums = adjacency @ values[:, :, None]  # Shape (N, M, 1)
        # Remove the last singleton dimension,
        value_sums = value_sums.squeeze(-1)  # Shape (N, M)

        return np.max(value_sums, axis=1)  # Shape (N,)


# reused in reflectance.py
def apply_plane_filters(values, theta, zv, elevation=None, azimuth=None):
    """apply angular view filters to a plane or point"""
    if not zv.is_volume():
        # apply normals/directions — shared weighting (same code, different theta)
        if zv.use_normal:
            values[theta > np.pi / 2] = 0

        # apply vertical/horizontal irradiances
        if zv.vert:
            values *= np.sin(theta)
        if zv.horiz:
            values *= abs(np.cos(theta))

        # FOV — branches on mode
        if elevation is not None and azimuth is not None:
            # View mode: rectangular FOV centered on view direction
            if zv.fov_vert < 180:
                values[np.abs(elevation) > np.radians(zv.fov_vert / 2)] = 0
            if zv.fov_horiz < 360:
                values[np.abs(azimuth) > np.radians(zv.fov_horiz / 2)] = 0
        else:
            # Normal mode: FOV centered on horizontal (existing)
            if zv.fov_vert < 180:
                values[theta < (np.pi / 2 - np.radians(zv.fov_vert / 2))] = 0
                values[theta > (np.pi / 2 + np.radians(zv.fov_vert / 2))] = 0
            # fov_horiz handled at aggregation time (worst-case search)
        return values
    else:
        return values


def _compute_view_angles(lamp, zv):
    """Compute theta, elevation, and azimuth for view mode.

    Returns (theta, elevation, azimuth) arrays, all shape (N,).
    theta: angle between incoming light and view direction
    elevation: angle above/below the view plane
    azimuth: angle left/right of the view direction
    """
    view_normals = zv.compute_view_normals()  # (N, 3)
    lamp_pos = np.asarray(lamp.surface.position, dtype="float64")

    # incoming = direction light arrives from (lamp → point), normalized
    incoming = zv.coords - lamp_pos
    dist = np.linalg.norm(incoming, axis=1, keepdims=True)
    dist = np.where(dist < 1e-12, 1.0, dist)
    incoming = incoming / dist

    # theta = angle between incoming light and view direction
    # light arriving FROM the lamp means we negate incoming to get "light direction toward eye"
    cos_theta = np.sum(-incoming * view_normals, axis=1)
    cos_theta = np.clip(cos_theta, -1.0, 1.0)
    theta = np.arccos(cos_theta)

    # Build per-point view frame: forward, right, up
    forward = view_normals  # (N, 3)
    world_up = np.array([0.0, 0.0, 1.0])

    right = np.cross(forward, world_up)
    right_norm = np.linalg.norm(right, axis=1, keepdims=True)
    # Fallback for forward parallel to world_up
    parallel = (right_norm.ravel() < 1e-6)
    if np.any(parallel):
        fallback_up = np.array([0.0, 1.0, 0.0])
        right[parallel] = np.cross(forward[parallel], fallback_up)
        right_norm[parallel] = np.linalg.norm(right[parallel], axis=1, keepdims=True)
    right = right / right_norm

    up = np.cross(right, forward)  # already normalized (cross of two unit vectors at 90°)

    # Project -incoming into view frame
    neg_incoming = -incoming
    fwd_component = np.sum(neg_incoming * forward, axis=1)
    right_component = np.sum(neg_incoming * right, axis=1)
    up_component = np.sum(neg_incoming * up, axis=1)

    # Elevation: angle above/below the forward-right plane
    horiz_dist = np.sqrt(fwd_component ** 2 + right_component ** 2)
    elevation = np.arctan2(up_component, horiz_dist)

    # Azimuth: angle left/right of forward in the forward-right plane
    azimuth = np.arctan2(right_component, fwd_component)

    return theta, elevation, azimuth
