import json, time
import numpy as np
from scipy.optimize import minimize

SEED = 7
RNG = np.random.default_rng(SEED)

def entropy(A, pi):
    x = np.clip(A, 1e-9, 1 - 1e-9)
    return float(np.sum(pi[:, None] * pi[None, :] * (x*np.log(x) + (1-x)*np.log(1-x))))

def edge_density(A, pi):
    return float(np.einsum('i,j,ij->', pi, pi, A))

def triangle_density(A, pi):
    return float(np.einsum('i,j,k,ij,jk,ki->', pi, pi, pi, A, A, A))

def wedge_density(A, pi):
    return float(np.einsum('i,j,k,ij,jk->', pi, pi, pi, A, A))

def optimize_entropy(pi, edge_target, triangle_target=None, seed=0):
    m = len(pi); rng = np.random.default_rng(seed)
    x0 = np.clip(edge_target + rng.normal(0, .04, (m, m)), .03, .97).ravel()
    constraints = [{'type':'eq', 'fun': lambda x: edge_density(x.reshape(m,m), pi)-edge_target}]
    if triangle_target is not None:
        constraints.append({'type':'eq', 'fun': lambda x: triangle_density(x.reshape(m,m), pi)-triangle_target})
    result = minimize(lambda x: entropy(x.reshape(m,m), pi), x0, method='SLSQP',
                      bounds=[(1e-5, .99999)]*(m*m), constraints=constraints,
                      options={'ftol':1e-12, 'maxiter':1000})
    return result.x.reshape(m,m), result

def exact_vs_sampled_motif():
    pi = np.array([.2, .5, .3])
    A = np.array([[.15,.4,.7],[.25,.6,.35],[.8,.3,.55]])
    exact = {'edge': edge_density(A,pi), 'wedge': wedge_density(A,pi), 'triangle': triangle_density(A,pi)}
    out = []
    for n in [100, 300, 1000, 3000, 10000, 30000]:
        errs = []
        for rep in range(40):
            z = RNG.choice(3, size=(n,3), p=pi)
            errs.append(np.mean(A[z[:,0],z[:,1]]*A[z[:,1],z[:,2]]*A[z[:,2],z[:,0]])-exact['triangle'])
        out.append({'samples':n, 'rmse':float(np.sqrt(np.mean(np.square(errs))))})
    return exact, out

def entropy_sweep():
    # Prediction 1: with only edge density, maximum entropy is homogeneous A=p.
    # Prediction 2: entropy decreases monotonically as a non-extremal target triangle
    # is forced away from its homogeneous value, while edge density remains fixed.
    pi = np.array([.25,.35,.40]); p=.30
    homogeneous_triangle=p**3
    rows=[]
    for tri in [homogeneous_triangle, .030, .040, .050, .060, .075, .090]:
        A,res=optimize_entropy(pi,p,tri,seed=11)
        rows.append({'target_triangle':tri, 'achieved_triangle':triangle_density(A,pi),
                     'edge_error':abs(edge_density(A,pi)-p), 'entropy':entropy(A,pi),
                     'heterogeneity':float(np.sqrt(np.sum(pi[:,None]*pi[None,:]*(A-p)**2))),
                     'success':bool(res.success)})
    one=[]
    for m in [2,3,4,5]:
        pi=np.ones(m)/m
        A,res=optimize_entropy(pi,p,None,seed=m)
        one.append({'m':m,'max_abs_deviation':float(np.max(np.abs(A-p))),
                    'weighted_rmse':float(np.sqrt(np.sum(pi[:,None]*pi[None,:]*(A-p)**2))),
                    'success':bool(res.success)})
    return {'homogeneous_triangle':homogeneous_triangle,'triangle_sweep':rows,'edge_only_block_sweep':one}

def propagation_benchmark():
    # Compare dense relation tensor application with the proposed block factorization.
    rng=np.random.default_rng(19); n=1800; m=8; r=3; d=32
    P=rng.dirichlet(np.ones(m), size=n).astype(np.float64)
    V=rng.normal(size=(n,d)).astype(np.float64)
    A=rng.uniform(.1,.9,size=(r,m,m)).astype(np.float64)
    # Dense kernel is n*n*r storage and application; use n=1800 to stay modest.
    W=np.einsum('vi,kij,wj->kvw',P,A,P)
    dense_bytes=W.nbytes
    t0=time.perf_counter(); dense=np.einsum('kvw,wd->vd',W,V); dense_time=time.perf_counter()-t0
    t0=time.perf_counter()
    S=np.einsum('wj,wd->jd',P,V)
    block=np.einsum('vi,kij,jd->vd',P,A,S)
    block_time=time.perf_counter()-t0
    rel_err=float(np.max(np.abs(dense-block))/(np.max(np.abs(dense))+1e-12))
    block_bytes=P.nbytes+A.nbytes
    return {'n':n,'m':m,'relations':r,'hidden':d,'dense_bytes':dense_bytes,
            'block_bytes':block_bytes,'memory_reduction':dense_bytes/block_bytes,
            'dense_seconds':dense_time,'block_seconds':block_time,'max_relative_error':rel_err}

def main():
    exact, sampling=exact_vs_sampled_motif()
    result={'seed':SEED,'motif_exact':exact,'motif_sampling_rmse':sampling,
            'entropy_sweep':entropy_sweep(),'propagation':propagation_benchmark()}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()