from dataclasses import dataclass
from typing import List, Union, Literal
from ortools.linear_solver import pywraplp
from curvaceous.curve import Curve

Number = Union[int, float]
Domain = Literal["x", "y"]


@dataclass
class Result:
    status: int
    xs: List[float]
    ys: List[float]


@dataclass 
class BaseConstraint:
    lb: Number
    ub: Number
    idx: List[int]

@dataclass
class XConstraint(BaseConstraint):
    lb: Number
    ub: Number
    domain: Literal["x"]

@dataclass
class YConstraint(BaseConstraint):
    lb: Number
    ub: Number
    domain: Literal["y"]
    name: str

Constraint = XConstraint | YConstraint


def maximize(curves: List[Curve], constraints: List[Constraint]) -> Result:
    solver, ws = _create_model(curves, constraints)
    status = solver.Solve()
    return _to_result(status, curves, ws)


def _create_model(curves: List[Curve], constraints: List[Constraint]):
    solver = pywraplp.Solver.CreateSolver("CBC")
    if not solver:
        raise RuntimeError("CBC solver unavailable")

    ws = []
    bs = []

    # first we define our costs and ys for constraints purposes
    costs = {i: [] for i, c in enumerate(constraints) if c.domain == "x"}
    ys_contrib = {i: [] for i, c in enumerate(constraints) if c.domain == "y"}

    objective_terms = []

    for idx, curve in enumerate(curves):
        k = len(curve.xs)

        # we assume the first y key is the thing we want to maximize
        y_names = list(curve.ys.keys())
        obj_y_name = y_names[0]
        obj_y = curve.ys[obj_y_name]

        # continuous weights for determining how much weight each segment gets
        w = [solver.NumVar(0.0, solver.infinity(), f"w_{idx}_{j}") for j in range(k)]
        ws.append(w)

        # binary selectors for choosing which segment is active
        b = [solver.IntVar(0, 1, f"b_{idx}_{j}") for j in range(k - 1)]
        bs.append(b)

        # make solution lie between points on the curve (for the math people: convex combination)
        solver.Add(solver.Sum(w[j] for j in range(k)) == 1)

        # linking constraints to make sure the model only interpolates between two adjacent points instead of across the whole curve
        solver.Add(w[0] <= b[0])
        for j in range(1, k - 1):
            solver.Add(w[j] <= b[j - 1] + b[j])
        solver.Add(w[k - 1] <= b[k - 2])

        solver.Add(solver.Sum(b[j] for j in range(k - 1)) == 1)

        # build the actual model and the local constraints (the ones that apply only to certain curves)
        for j in range(k):
            objective_terms.append(w[j] * float(obj_y[j]))

            for i, c in enumerate(constraints):
                if idx not in c.idx:
                    continue

                if c.domain == "x":
                    costs[i].append(w[j] * float(curve.xs[j]))

                elif c.domain == "y":
                    y_values = curve.ys[c.name]
                    ys_contrib[i].append(w[j] * float(y_values[j]))

    # add global constraints
    for i, c in enumerate(constraints):
        if c.domain == "x":
            expr = solver.Sum(costs[i])
        elif c.domain == "y":
            expr = solver.Sum(ys_contrib[i])
        else:
            continue

        if c.lb is not None:
            solver.Add(expr >= c.lb)
        if c.ub is not None:
            solver.Add(expr <= c.ub)

    solver.Maximize(solver.Sum(objective_terms))

    return solver, ws


def _to_result(status, curves, ws):
    if status == pywraplp.Solver.OPTIMAL:
        return Result(status, *_compute_xs_and_ys(curves, ws))
    return Result(status, None, None)


def _compute_xs_and_ys(curves, ws):
    xs = []
    ys = []

    for i, curve in enumerate(curves):
        k = len(curve.xs)

        x_val = sum(
            ws[i][j].solution_value() * float(curve.xs[j])
            for j in range(k)
        )
        xs.append(x_val)

        y_vals = {}
        for name, y_arr in curve.ys.items():
            y_vals[name] = sum(
                ws[i][j].solution_value() * float(y_arr[j])
                for j in range(k)
            )

        ys.append(y_vals)

    return xs, ys
