import json
import numpy as np

SEED = 1099
np.random.seed(SEED)

def K(x, y, ell=0.18):
    x = np.asarray(x)[:, None]; y = np.asarray(y)[None, :]
    return np.exp(-0.5*((x-y)/ell)**2)

def var(train, z, noise=0.04):
    train = np.asarray(train, float); z = np.asarray(z, float)
    A = K(train, train) + noise*np.eye(len(train))
    cross = K(z, train)
    q = np.linalg.solve(A, cross.T)
    return np.maximum(0., 1.-np.sum(cross*q.T, axis=1))

def red(train, action, z, w):
    return float(np.dot(w, var(train,z)-var(np.r_[train,action],z)))

def controlled_overlap():
    train=np.array([-1.,1.]); z=np.linspace(-1,1,501); w=np.ones(len(z))/len(z)
    observed=np.linspace(-.12,.12,8)
    remote=np.linspace(.68,.88,8)
    rows=[]
    # Number of repeated/nearby points increases; all action sizes remain fixed.
    for m in range(9):
        b=np.r_[observed[:m], remote[:8-m]]
        rows.append({'near_fraction':m/8., 'reduction_after_first':red(np.r_[train,observed],b,z,w)})
    return rows

def affine_tree():
    train=np.array([-1.,1.]); z=np.linspace(-1,1,301); w=np.ones(len(z))/len(z)
    centers=np.array([-.72,-.42,-.10,.22,.55,.82])
    acts=[np.linspace(c-.1,c+.1,5) for c in centers]
    gamma=.9; n=len(acts); lines={}
    for i in range(n):
        first=red(train,acts[i],z,w)
        for j in range(n):
            second=red(np.r_[train,acts[i]],acts[j],z,w)
            lines[i,j]=(first+gamma*second, gamma*abs(i-j)/(n-1))
    def policy(lam):
        best=[]
        for i in range(n):
            j=max(range(n),key=lambda q: lines[i,q][0]-lam*lines[i,q][1])
            best.append((lines[i,j][0]-lam*lines[i,j][1],(i,j)))
        return max(best)
    grid=np.linspace(0,2,2001)
    policies=[policy(x)[1] for x in grid]
    # Locate the first change whose adjacent paths have distinct costs.
    chosen=None
    for k in range(len(grid)-1):
        if policies[k]!=policies[k+1]:
            p,q=policies[k],policies[k+1]
            if abs(lines[p][1]-lines[q][1])>1e-10:
                A,B=lines[p]; C,D=lines[q]
                exact=(A-C)/(B-D)
                chosen={'path_before':p,'path_after':q,'predicted_lambda':exact,
                        'observed_grid_midpoint':(grid[k]+grid[k+1])/2,
                        'grid_error':abs(exact-(grid[k]+grid[k+1])/2)}
                break
    return chosen

def comparison():
    train0=np.array([-1.,1.]); z=np.linspace(-1,1,401); w=np.ones(len(z))/len(z)
    centers=np.array([-.72,-.42,-.10,.22,.55,.82])
    acts=[np.linspace(c-.1,c+.1,5) for c in centers]; gamma=.9
    def tree(tr,lam=.8):
        vals=[]
        for i,a in enumerate(acts):
            r=red(tr,a,z,w); future=max(red(np.r_[tr,a],b,z,w)-lam*abs(i-j)/(len(acts)-1) for j,b in enumerate(acts))
            vals.append(r+gamma*future)
        return int(np.argmax(vals))
    out={}
    for method in ('random','greedy','tree'):
        tr=train0.copy(); total=0.; chosen=[]
        for t in range(3):
            if method=='random': i=np.random.default_rng(SEED+t).integers(len(acts))
            elif method=='greedy': i=int(np.argmax([red(tr,a,z,w) for a in acts]))
            else: i=tree(tr)
            total+=red(tr,acts[i],z,w); tr=np.r_[tr,acts[i]]; chosen.append(int(i))
        out[method]={'cumulative_reduction':total,'actions':chosen,'final_mean_variance':float(np.mean(var(tr,z)))}
    return out

def main():
    z=np.linspace(-1,1,401); train=np.array([-.85,-.25,.35,.9])
    before=var(train,z); after=var(np.r_[train,0.],z)
    max_inc=float(np.max(after-before))
    overlap=controlled_overlap()
    vals=[r['reduction_after_first'] for r in overlap]
    # Prediction: variance reduction decreases as the fraction of redundant points rises.
    mono=all(vals[i+1] <= vals[i]+1e-9 for i in range(len(vals)-1))
    tree=affine_tree()
    result={'seed':SEED,
      'prediction_1_variance_nonincrease':{'predicted_max_increase':0.,'observed_max_increase':max_inc,'confirmed':max_inc<=1e-9},
      'prediction_2_redundancy_diminishes_reward':{'predicted':'monotone decrease with near/overlap fraction','sweep':overlap,'confirmed':mono,'endpoint_ratio':vals[-1]/vals[0]},
      'prediction_3_cost_crossover':{'predicted':'lambda=(A-C)/(B-D) for affine path values','result':tree,'confirmed':tree is not None and tree['grid_error']<.002},
      'comparison':comparison()}
    with open('results.json','w') as f: json.dump(result,f,indent=2,default=lambda x: x.item() if hasattr(x, 'item') else x)
    print(json.dumps(result,indent=2,default=lambda x: x.item() if hasattr(x, 'item') else x))
if __name__=='__main__': main()