from dataclasses import dataclass
from typing import List, Union, Literal
import pulp
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 Constraint:
    lb: Number
    ub: Number
    idx: List[int]
    domain: Domain



def maximize(curves: List[Curve], constraints: List[Constraint]) -> Result:
    m, ws = _create_model(curves, constraints)
    status = m.solve(pulp.PULP_CBC_CMD(msg=False))
    return _to_result(status, curves, ws)


def _create_model(
    curves: List[Curve],
    constraints: List[Constraint],
    print_logs: bool = False,
):
    # Create maximization model
    m = pulp.LpProblem("maximize_curves", pulp.LpMaximize)

    # Objective expression
    value = pulp.lpSum([])

    bs = []
    ws = []

    # Expressions for constraints
    costs = {
        i: pulp.lpSum([])
        for i, c in enumerate(constraints)
        if c.domain == "x"
    }
    ys = {
        i: pulp.lpSum([])
        for i, c in enumerate(constraints)
        if c.domain == "y"
    }

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

        # Continuous weights
        w = [
            pulp.LpVariable(f"w_{idx}_{j}", lowBound=0, cat="Continuous")
            for j in range(k)
        ]

        # Binary selectors
        b = [
            pulp.LpVariable(f"b_{idx}_{j}", cat="Binary")
            for j in range(k - 1)
        ]

        ws.append(w)
        bs.append(b)

        # Sum of weights equals 1
        m += pulp.lpSum(w[j] for j in range(k)) == 1

        # Linking constraints
        m += w[0] <= b[0]
        for j in range(1, k - 1):
            m += w[j] <= b[j - 1] + b[j]
        m += w[k - 1] <= b[k - 2]

        # Exactly one segment active
        m += pulp.lpSum(b[j] for j in range(k - 1)) == 1

        for j in range(k):
            # Objective contribution
            value += w[j] * float(curve.ys[j])

            # Constraint contributions
            for i, c in enumerate(constraints):
                if idx in c.idx:
                    if c.domain == "x":
                        costs[i] += w[j] * float(curve.xs[j])
                    elif c.domain == "y":
                        ys[i] += w[j] * float(curve.ys[j])

    # Add constraints
    for i, c in enumerate(constraints):
        if c.domain == "x":
            expr = costs[i]
        elif c.domain == "y":
            expr = ys[i]
        else:
            continue

        if c.lb is not None:
            m += expr >= c.lb
        if c.ub is not None:
            m += expr <= c.ub

    # Objective
    m += value

    return m, ws

def _to_result(status, curves, ws):
    if status == pulp.LpStatusOptimal:
        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.append(sum(ws[i][j].value() * float(curve.xs[j]) for j in range(0, k)))
        ys.append(sum(ws[i][j].value() * float(curve.ys[j]) for j in range(0, k)))
    return (xs, ys)
