import json, math, random
import numpy as np
from sklearn.linear_model import Lasso
from sklearn.metrics import mean_squared_error

SEED = 3053
rng = np.random.default_rng(SEED)


def solve_lasso(A, b, lam, alpha=None, max_iter=20000):
    # sklearn objective is (1/(2n))*||Aw-b||^2 + alpha ||w||_1
    if alpha is None: alpha = lam / len(b)
    model = Lasso(alpha=alpha, fit_intercept=False, max_iter=max_iter,
                  tol=1e-11, selection='cyclic', random_state=SEED)
    model.fit(A, b)
    w = model.coef_.copy()
    resid = A @ w - b
    primal = .5 * np.dot(resid, resid) + lam * np.abs(w).sum()
    # Dual y convention: y = Aw-b; maximize -.5||y||^2 - b^T y,
    # constrained by ||A^T y||inf <= lambda.
    yraw = resid.copy()
    aty = A.T @ yraw
    scale = min(1.0, lam / (np.max(np.abs(aty)) + 1e-30))
    y = yraw * scale
    dual = -.5*np.dot(y,y) - np.dot(b,y)
    gap = max(0.0, primal-dual)
    return w, y, primal, dual, gap


def core_math_check():
    n, d = 80, 24
    A = rng.normal(size=(n,d))
    # Correlated response with deliberately many zero optimum coordinates.
    truth = np.zeros(d); truth[[1,5,9,14]] = [2.0,-1.5,.8,1.2]
    b = A @ truth + .15*rng.normal(size=n)
    lam = 2.0
    w,y,pr,du,gap = solve_lasso(A,b,lam)
    aty = A.T @ y
    R = math.sqrt(2*gap + 1e-18) # dual objective is 1-strongly concave
    safe = np.abs(aty) + R*np.linalg.norm(A,axis=0) < lam
    zeros = np.abs(w) < 2e-6
    no_false_safe = bool(np.all(~safe | zeros))
    # At optimum, KKT complementarity: nonzero coordinates have |A'y|=lambda.
    kkt_nonzero = float(np.max(np.abs(np.abs(aty[np.abs(w)>2e-6])-lam))) if np.any(np.abs(w)>2e-6) else 0.
    # Equivariance: selecting columns then dualizing is exactly the same
    # reduced matrix as row/sample selection then dualizing.
    fmask = np.abs(aty) >= np.quantile(np.abs(aty), .5)
    smask = np.arange(n) % 2 == 0
    reduced_both = A[smask][:,fmask]
    dualized_after_feature = reduced_both
    feature_after_dual = A[smask][:,fmask]
    equiv_err = float(np.max(np.abs(dualized_after_feature-feature_after_dual)))
    # Show safety becomes useful as gap shrinks using exact-ish solution.
    caught = int(np.sum(safe)); zero_count=int(np.sum(zeros))
    return dict(n=n,d=d, primal=pr,dual=du,gap=gap,radius=R,
                safe_features=caught, optimizer_zero_features=zero_count,
                no_false_safe=no_false_safe,kkt_nonzero_error=kkt_nonzero,
                mask_equivariance_max_error=equiv_err)


def fit_on_masks(Atr, btr, Aval, bval, fmask, smask, lam):
    cols=np.flatnonzero(fmask); rows=np.flatnonzero(smask)
    if len(cols)==0 or len(rows)==0: return (float('inf'), float('inf'), 0)
    w,y,pr,du,gap=solve_lasso(Atr[np.ix_(rows,cols)], btr[rows], lam)
    pred=Aval[:,cols]@w
    return mean_squared_error(bval,pred), pr, len(cols)


def mini_experiment():
    # Feature pruning is tested on a held-out set; sample pruning only changes
    # training rows. Dual sample score is the separable conjugate excess score.
    ntr,nva,d=240,160,60
    A=rng.normal(size=(ntr+nva,d)); truth=np.zeros(d)
    truth[[2,7,13,31,44]]=[2.2,-1.7,1.1,.9,-1.3]
    b=A@truth + .35*rng.normal(size=ntr+nva)
    Atr,Ava=A[:ntr],A[ntr:]; btr,bva=b[:ntr],b[ntr:]
    lam=1.0
    fullmask=np.ones(d,dtype=bool)
    # Obtain dual scores from full training solution.
    w,y,pr,du,gap=solve_lasso(Atr,btr,lam); aty=Atr.T@y
    R=math.sqrt(2*gap+1e-18)
    safe=np.abs(aty)+R*np.linalg.norm(Atr,axis=0)<lam
    # Explicit heuristic fallback ranks by dual correlation, as specified.
    dual_feat_order=np.argsort(np.abs(aty))
    mag_order=np.argsort(np.abs(w))
    # squared-loss f_i^* excess over its minimum = .5*(A_i w)^2
    sample_score=.5*(Atr@w)**2
    dual_sample_order=np.argsort(sample_score) # low contribution first
    random_order=rng.permutation(ntr)
    results=[]
    for frac in [0.25,0.50,0.75]:
        kf=max(1,int(d*(1-frac))); ks=max(1,int(ntr*(1-frac)))
        # Baseline: magnitude features + random samples.
        bm=np.ones(d,bool); bm[mag_order[:d-kf]]=False
        bs=np.ones(ntr,bool); bs[random_order[:ntr-ks]]=False
        # Idea: permanently safe-remove features, then dual-correlation rank
        # remaining features if target exceeds safe removals; dual sample scores.
        im=np.ones(d,bool); im[safe]=False
        if im.sum()>kf:
            remaining=np.flatnonzero(im)
            im[remaining[np.argsort(np.abs(aty[remaining]))[:len(remaining)-kf]]]=False
        elif im.sum()<kf:
            # conservative rule cannot retain fewer than requested; leave mask safe
            pass
        isamp=np.ones(ntr,bool); isamp[dual_sample_order[:ntr-ks]]=False
        bval_m,_,bd=fit_on_masks(Atr,btr,Ava,bva,bm,bs,lam)
        ival_m,_,idd=fit_on_masks(Atr,btr,Ava,bva,im,isamp,lam)
        results.append(dict(reduction=frac,baseline_val_mse=bval_m,idea_val_mse=ival_m,
                            baseline_features=int(bm.sum()),idea_features=int(im.sum()),
                            baseline_samples=int(bs.sum()),idea_samples=int(isamp.sum()),
                            safe_removed=int(safe.sum())))
    return dict(train=ntr,validation=nva,features=d,full_gap=gap,
                full_nonzero=int(np.sum(np.abs(w)>1e-6)),results=results)


def main():
    math_result=core_math_check()
    exp=mini_experiment()
    out={'seed':SEED,'math_check':math_result,'mini_experiment':exp}
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()