"""Alternative scatterer geometries for Helmholtz FEM problems.

Supports non-circular dielectric scatterers: square, multiple circles, etc.
Each class overrides only the geometry-dependent methods of HelmholtzProblem.
"""

from typing import Any, Dict, Union

import numpy as np
from skfem import Mesh

from environment.helmholtz import (
    HelmholtzProblem,
    _compute_residual_indicator,
)


# ═══════════════════════════════════════════════════════════════════
#  Square dielectric scatterer
# ═══════════════════════════════════════════════════════════════════

class HelmholtzProblemSquare(HelmholtzProblem):
    """Helmholtz problem with a square dielectric scatterer.

    Extra config keys under helmholtz.scatterer.square:
        half_side: float  — half side length (default 0.2)
        angle: float      — rotation in radians (default 0.0)
    """

    def __init__(
        self,
        *,
        fem_config: Dict[Union[str, int], Any],
        random_state: np.random.RandomState = np.random.RandomState(),
    ):
        sc = fem_config.get("helmholtz", {}).get("scatterer", {})
        sq = sc.get("square", {})
        self._sq_cx = float(sq.get("cx", sc.get("cx", 0.5)))
        self._sq_cy = float(sq.get("cy", sc.get("cy", 0.5)))
        self._sq_half = float(sq.get("half_side", sc.get("radius", 0.2)))
        self._sq_angle = float(sq.get("angle", 0.0))
        self._sq_eps_r = float(sc.get("eps_r", 2.0))

        super().__init__(fem_config=fem_config, random_state=random_state)
        self._eps_r = self._sq_eps_r

    # ── geometry helpers ──

    def _rotate_xy(self, x, y):
        """Rotate coordinates back to scatterer-local frame."""
        if self._sq_angle == 0:
            return x - self._sq_cx, y - self._sq_cy
        c, s = np.cos(-self._sq_angle), np.sin(-self._sq_angle)
        dx, dy = x - self._sq_cx, y - self._sq_cy
        return c * dx - s * dy, s * dx + c * dy

    def _in_square(self, x, y):
        xr, yr = self._rotate_xy(x, y)
        return (np.abs(xr) <= self._sq_half) & (np.abs(yr) <= self._sq_half)

    # ── FEM assembly (called at quadrature points) ──

    def _eps_r_at_quad_points(self, x, y):
        return np.where(self._in_square(x, y), self._sq_eps_r, 1.0)

    # ── midpoint eps_r for error estimation / features ──

    def eps_r_at_midpoints(self, mesh: Mesh) -> np.ndarray:
        pts = np.mean(mesh.p[:, mesh.t], axis=1).T
        return np.where(self._in_square(pts[:, 0], pts[:, 1]), self._sq_eps_r, 1.0)

    # ── override error estimation ──

    def get_error_estimate_per_element(self, basis, solution):
        eps_r_arr = self.eps_r_at_midpoints(basis.mesh)
        return {"indicator": _compute_residual_indicator(
            basis.mesh, solution, k=self._k, eps_r=eps_r_arr)}

    # ── override features ──

    def element_features(self, mesh, element_feature_names):
        features_list = []
        if "epsilon_r" in element_feature_names:
            features_list.append(self.eps_r_at_midpoints(mesh)[:, None])
        return np.concatenate(features_list, axis=1) if features_list else None

    # ── Nyquist enforcement uses square bounding box ──

    def _enforce_nyquist_in_dielectric(self, mesh, N=1.5, max_iter=10):
        lambda_d = 2.0 * np.pi / (self._k * np.sqrt(self._sq_eps_r))
        h_max = lambda_d / N

        for _ in range(max_iter):
            i0, i1, i2 = mesh.t[0], mesh.t[1], mesh.t[2]
            x0, y0 = mesh.p[0, i0], mesh.p[1, i0]
            x1, y1 = mesh.p[0, i1], mesh.p[1, i1]
            x2, y2 = mesh.p[0, i2], mesh.p[1, i2]

            e01 = np.sqrt((x1 - x0) ** 2 + (y1 - y0) ** 2)
            e12 = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
            e20 = np.sqrt((x0 - x2) ** 2 + (y0 - y2) ** 2)
            h_K = np.maximum(np.maximum(e01, e12), e20)

            midpoints = np.mean(mesh.p[:, mesh.t], axis=1).T
            in_dielectric = self._in_square(midpoints[:, 0], midpoints[:, 1])
            to_refine = np.where(in_dielectric & (h_K > h_max))[0]
            if len(to_refine) == 0:
                break
            mesh = mesh.refined(to_refine)
        return mesh

    # ── visualization overlay ──

    def additional_plots_from_mesh(self, mesh: Mesh) -> Dict:
        corners = np.array([
            [-self._sq_half, -self._sq_half],
            [ self._sq_half, -self._sq_half],
            [ self._sq_half,  self._sq_half],
            [-self._sq_half,  self._sq_half],
            [-self._sq_half, -self._sq_half],
        ])
        if self._sq_angle != 0:
            c, s = np.cos(self._sq_angle), np.sin(self._sq_angle)
            rot = np.array([[c, -s], [s, c]])
            corners = corners @ rot.T
        corners[:, 0] += self._sq_cx
        corners[:, 1] += self._sq_cy
        return {"square_outline": (corners[:, 0], corners[:, 1])}


# ═══════════════════════════════════════════════════════════════════
#  Multi-circle dielectric scatterer
# ═══════════════════════════════════════════════════════════════════

class HelmholtzProblemMultiCircle(HelmholtzProblem):
    """Helmholtz problem with multiple circular dielectric scatterers.

    Extra config key under helmholtz.scatterer:
        circles: list of dicts, each with cx, cy, radius, eps_r
    """

    def __init__(
        self,
        *,
        fem_config: Dict[Union[str, int], Any],
        random_state: np.random.RandomState = np.random.RandomState(),
    ):
        sc = fem_config.get("helmholtz", {}).get("scatterer", {})
        circles_cfg = sc.get("circles", None)
        if circles_cfg is None:
            circles_cfg = [{
                "cx": sc.get("cx", 0.35), "cy": sc.get("cy", 0.35),
                "radius": sc.get("radius", 0.12), "eps_r": sc.get("eps_r", 3.0),
            }, {
                "cx": sc.get("cx", 0.65) if "cx2" not in sc else sc["cx2"], "cy": sc.get("cy", 0.65) if "cy2" not in sc else sc["cy2"],
                "radius": sc.get("radius", 0.12), "eps_r": sc.get("eps_r", 3.0),
            }]

        self._circles = []
        for c in circles_cfg:
            self._circles.append({
                "cx": float(c["cx"]),
                "cy": float(c["cy"]),
                "radius": float(c["radius"]),
                "eps_r": float(c.get("eps_r", 2.0)),
            })

        super().__init__(fem_config=fem_config, random_state=random_state)
        sc_primary = self._circles[0]
        self._eps_r = sc_primary["eps_r"]
        self._cx = sc_primary["cx"]
        self._cy = sc_primary["cy"]
        self._radius = sc_primary["radius"]

    # ── geometry ──

    def _eps_r_at_point(self, x, y):
        """Return eps_r at arbitrary points (broadcast-safe)."""
        out = np.ones_like(x, dtype=float)
        for c in self._circles:
            in_c = (x - c["cx"]) ** 2 + (y - c["cy"]) ** 2 <= c["radius"] ** 2
            out = np.where(in_c, c["eps_r"], out)
        return out

    # ── FEM assembly ──

    def _eps_r_at_quad_points(self, x, y):
        return self._eps_r_at_point(x, y)

    # ── midpoint eps_r ──

    def eps_r_at_midpoints(self, mesh: Mesh) -> np.ndarray:
        pts = np.mean(mesh.p[:, mesh.t], axis=1).T
        return self._eps_r_at_point(pts[:, 0], pts[:, 1])

    # ── error estimation ──

    def get_error_estimate_per_element(self, basis, solution):
        eps_r_arr = self.eps_r_at_midpoints(basis.mesh)
        return {"indicator": _compute_residual_indicator(
            basis.mesh, solution, k=self._k, eps_r=eps_r_arr)}

    # ── features ──

    def element_features(self, mesh, element_feature_names):
        features_list = []
        if "epsilon_r" in element_feature_names:
            features_list.append(self.eps_r_at_midpoints(mesh)[:, None])
        return np.concatenate(features_list, axis=1) if features_list else None

    # ── Nyquist enforcement for all circles ──

    def _enforce_nyquist_in_dielectric(self, mesh, N=1.5, max_iter=15):
        for _ in range(max_iter):
            i0, i1, i2 = mesh.t[0], mesh.t[1], mesh.t[2]
            x0, y0 = mesh.p[0, i0], mesh.p[1, i0]
            x1, y1 = mesh.p[0, i1], mesh.p[1, i1]
            x2, y2 = mesh.p[0, i2], mesh.p[1, i2]

            e01 = np.sqrt((x1 - x0) ** 2 + (y1 - y0) ** 2)
            e12 = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
            e20 = np.sqrt((x0 - x2) ** 2 + (y0 - y2) ** 2)
            h_K = np.maximum(np.maximum(e01, e12), e20)

            midpoints = np.mean(mesh.p[:, mesh.t], axis=1).T
            eps_r_at_mid = self._eps_r_at_point(midpoints[:, 0], midpoints[:, 1])
            lambda_local = 2.0 * np.pi / (self._k * np.sqrt(np.maximum(eps_r_at_mid, 1.0)))
            h_max = lambda_local / N

            to_refine = np.where((eps_r_at_mid > 1.0) & (h_K > h_max))[0]
            if len(to_refine) == 0:
                break
            mesh = mesh.refined(to_refine)
        return mesh

    # ── visualization overlay ──

    def additional_plots_from_mesh(self, mesh: Mesh) -> Dict:
        result = {}
        for i, c in enumerate(self._circles):
            theta = np.linspace(0, 2 * np.pi, 128)
            result[f"circle_{i}"] = (
                c["cx"] + c["radius"] * np.cos(theta),
                c["cy"] + c["radius"] * np.sin(theta),
            )
        return result


# ═══════════════════════════════════════════════════════════════════
#  Factory functions  (mirror create_helmholtz_problem)
# ═══════════════════════════════════════════════════════════════════

def create_helmholtz_problem_square(
    *, fem_config: Dict[Union[str, int], Any],
    random_state: np.random.RandomState = np.random.RandomState(),
) -> HelmholtzProblemSquare:
    return HelmholtzProblemSquare(fem_config=fem_config, random_state=random_state)


def create_helmholtz_problem_multi_circle(
    *, fem_config: Dict[Union[str, int], Any],
    random_state: np.random.RandomState = np.random.RandomState(),
) -> HelmholtzProblemMultiCircle:
    return HelmholtzProblemMultiCircle(fem_config=fem_config, random_state=random_state)
