from __future__ import annotations

import math
from typing import Iterable, Optional, Sequence, Union

import numpy as np
import cvxpy as cp


ArrayLike = Union[np.ndarray, Sequence[Sequence[float]]]


def TROP_TWFE_average(
    Y: ArrayLike,
    W: ArrayLike,
    treated_units: Sequence[int],
    lambda_unit: float,
    lambda_time: float,
    lambda_nn: float,
    treated_periods: int = 10,
    solver: Optional[str] = None,
    verbose: bool = False,
) -> float:
    """
    Triply Robust Panel (TROP) estimator in a TWFE framework with:
      - distance-based unit weights (lambda_unit)
      - distance-based time weights (lambda_time)
      - optional low-rank regression adjustment (nuclear norm penalty, lambda_nn)

    Parameters
    ----------
    Y:
        Outcome matrix of shape (N, T).
    W:
        Treatment indicator matrix of shape (N, T). Typically binary {0,1}.
        Convention in the provided codebase: treated units are treated in the final
        `treated_periods` columns, but the function will run for any W.
    treated_units:
        Indices of treated units (row indices into Y/W).
    lambda_unit:
        Nonnegative tuning parameter controlling unit-weight decay.
    lambda_time:
        Nonnegative tuning parameter controlling time-weight decay.
    lambda_nn:
        Nuclear norm penalty weight for the low-rank component L.
        Use np.inf to disable low-rank adjustment (i.e., no L term).
    treated_periods:
        Number of treated (post) periods at the end of the panel used to define:
          - the time-distance center
          - the pre-period mask for unit distance computation
    solver:
        Optional CVXPY solver name. If None, chooses a reasonable default:
          - if lambda_nn is finite: uses SCS (supports nuclear norm / SDP forms)
          - if lambda_nn is infinite: uses OSQP (fast for pure quadratic problems)
    verbose:
        Passed to CVXPY solve().

    Returns
    -------
    float
        Estimated average treatment effect tau.

    """
    Y = np.asarray(Y, dtype=float)
    W = np.asarray(W, dtype=float)

    if Y.ndim != 2 or W.ndim != 2:
        raise ValueError(f"Y and W must be 2D arrays. Got Y.ndim={Y.ndim}, W.ndim={W.ndim}.")
    if Y.shape != W.shape:
        raise ValueError(f"Y and W must have the same shape. Got Y={Y.shape}, W={W.shape}.")

    N, T = Y.shape

    if not isinstance(treated_periods, int) or treated_periods <= 0:
        raise ValueError("treated_periods must be a positive integer.")
    if treated_periods >= T:
        raise ValueError(f"treated_periods must be < T. Got treated_periods={treated_periods}, T={T}.")

    treated_units_arr = np.asarray(treated_units, dtype=int)
    if treated_units_arr.size == 0:
        raise ValueError("treated_units must contain at least one unit index.")
    if np.any(treated_units_arr < 0) or np.any(treated_units_arr >= N):
        raise ValueError(f"treated_units contains out-of-range indices for N={N}: {treated_units_arr}")

    if lambda_unit < 0 or lambda_time < 0:
        raise ValueError("lambda_unit and lambda_time should be nonnegative.")

    # ---------------------------------------------------------------------
    # Distance-based time weights
    # ---------------------------------------------------------------------
    # Distance to the center of the treated block near the end of the panel.
    # dist_time = abs(arange(T) - (T - treated_periods/2))
    center = T - treated_periods / 2.0
    dist_time = np.abs(np.arange(T, dtype=float) - center)

    # ---------------------------------------------------------------------
    # Distance-based unit weights
    # ---------------------------------------------------------------------
    average_treated = np.mean(Y[treated_units_arr, :], axis=0)

    # Pre-period mask: 1 in pre, 0 in treated/post
    mask = np.ones((N, T), dtype=float)
    mask[:, -treated_periods:] = 0.0

    # RMS distance to average treated trajectory over pre-periods
    # dist_unit[i] = sqrt( sum_pre (avg_tr - Y_i)^2 / (#pre) )
    A = np.sum(((average_treated - Y) ** 2) * mask, axis=1)
    B = np.sum(mask, axis=1)

    if np.any(B == 0):
        raise ValueError(
            "Pre-period mask has zero pre-periods for at least one unit."
        )

    dist_unit = np.sqrt(A / B)

    # Convert distances to weights
    delta_unit = np.exp(-lambda_unit * dist_unit)          # shape (N,)
    delta_time = np.exp(-lambda_time * dist_time)          # shape (T,)
    delta = np.outer(delta_unit, delta_time)               # shape (N, T)

    # ---------------------------------------------------------------------
    # CVXPY problem: weighted TWFE
    # ---------------------------------------------------------------------
    unit_effects = cp.Variable((1, N))
    time_effects = cp.Variable((1, T))
    mu = cp.Variable()     # intercept
    tau = cp.Variable()    # treatment effect

    # Broadcast TWFE components to N x T
    unit_factor = cp.kron(np.ones((T, 1)), unit_effects).T
    time_factor = cp.kron(np.ones((N, 1)), time_effects)

    is_low_rank = not math.isinf(float(lambda_nn))

    if is_low_rank:
        L = cp.Variable((N, T))
        residual = Y - mu - unit_factor - time_factor - L - W * tau
        loss = cp.sum_squares(cp.multiply(residual, delta)) + float(lambda_nn) * cp.norm(L, "nuc")
        default_solver = "SCS"  # robust choice for nuclear norm problems
    else:
        residual = Y - mu - unit_factor - time_factor - W * tau
        loss = cp.sum_squares(cp.multiply(residual, delta))
        default_solver = "OSQP"  # fast for pure quadratic objective

    prob = cp.Problem(cp.Minimize(loss))

    chosen_solver = solver or default_solver
    prob.solve(solver=chosen_solver, verbose=verbose)

    if tau.value is None or not np.isfinite(tau.value):
        raise RuntimeError(
            "Optimization did not return a valid tau. "
            f"Solver={chosen_solver}, status={prob.status}."
        )

    return float(tau.value)