import torch
import numpy as np
from .moo import cagrad_update_k2, raco_update_k2

class ConflictingQuadratics:
    def __init__(self, dim=100, conflict_angle=np.pi * 0.6, seed=42):
        rng = np.random.RandomState(seed)
        self.dim = dim
        self.A1 = rng.randn(dim, dim).astype(np.float32)
        self.A1 = self.A1.T @ self.A1 / dim
        self.b1 = rng.randn(dim).astype(np.float32) * 0.1

        self.A2 = rng.randn(dim, dim).astype(np.float32)
        self.A2 = self.A2.T @ self.A2 / dim
        self.b2 = rng.randn(dim).astype(np.float32) * 0.1

        v = rng.randn(dim).astype(np.float32)
        v = v / np.linalg.norm(v)
        u = rng.randn(dim).astype(np.float32)
        u = u - np.dot(u, v) * v
        u = u / np.linalg.norm(u)
        g1_0 = v
        g2_0 = v * np.cos(conflict_angle) + u * np.sin(conflict_angle)
        scale = np.linalg.norm(self.A1 @ self.b1 + self.b1)
        self.b1 = self.b1 - g1_0 * scale
        self.b2 = self.b2 - g2_0 * scale

def run_synthetic_comparison(dim=100, conflict_angle=np.pi*0.6, steps=200, lr=0.01, seed=42):
    import math
    prob = ConflictingQuadratics(dim=dim, conflict_angle=conflict_angle, seed=seed)
    A1_t = torch.from_numpy(prob.A1).float()
    b1_t = torch.from_numpy(prob.b1).float()
    A2_t = torch.from_numpy(prob.A2).float()
    b2_t = torch.from_numpy(prob.b2).float()
    theta0 = torch.from_numpy(np.random.RandomState(seed + 1).randn(dim).astype(np.float32) * 0.5)

    methods = {}
    for name, use_clip, w1, w2 in [
        ("DPO_LW", False, 0.7, 0.3),
        ("CAGrad", False, 0.7, 0.3),
        ("RACO", True, 0.7, 0.3),
    ]:
        theta = theta0.clone().requires_grad_(True)
        losses = {0: [], 1: []}
        opt = torch.optim.SGD([theta], lr=lr)
        for t in range(steps):
            opt.zero_grad()
            l1 = 0.5 * (theta @ A1_t @ theta) + b1_t @ theta
            l2 = 0.5 * (theta @ A2_t @ theta) + b2_t @ theta

            g1 = torch.autograd.grad(l1, theta, retain_graph=True)[0]
            g2 = torch.autograd.grad(l2, theta, retain_graph=True)[0]

            g1_list = [g1]
            g2_list = [g2]

            if name == "DPO_LW":
                g_update = g1 * w1 + g2 * w2
            elif name == "CAGrad":
                gl = cagrad_update_k2(g1_list, g2_list, w1, w2, c=0.4)
                g_update = gl[0]
            elif name == "RACO":
                gl = raco_update_k2(g1_list, g2_list, w1, w2, c=0.4)
                g_update = gl[0]

            theta.grad = g_update.detach() if hasattr(g_update, 'detach') else g_update
            opt.step()

            losses[0].append(float(l1.detach()))
            losses[1].append(float(l2.detach()))

        methods[name] = {"theta": theta.detach().numpy().copy(), "losses": losses}

    return methods