import json
from pathlib import Path
import numpy as np

SEED = 2304
EPS = 1e-12
rng = np.random.default_rng(SEED)


def clr(x):
    x = np.asarray(x, float)
    lx = np.log(x)
    return lx - np.mean(lx, axis=-1, keepdims=True)


def row(s, c):
    c = np.asarray(c, float)
    c = c / c.sum()
    return np.r_[s, (1.0 - s) * c]


def split(p):
    s = float(p[0])
    c = np.maximum(np.asarray(p[1:], float) / (1.0 - s), EPS)
    return s, c / c.sum()


def distances(a, b):
    sa, ca = split(a)
    sb, cb = split(b)
    q = np.log(ca) - np.log(cb)
    g = float(q.mean())
    z = q - g
    dc2 = float(np.mean(z * z))
    ua = np.log(sa / (1.0 - sa))
    ub = np.log(sb / (1.0 - sb))
    ds2 = 0.5 * (ua - ub) ** 2
    full = float(np.sum((clr(a) - clr(b)) ** 2))
    m = len(ca)
    D = m + 1
    # Exact CLR decomposition. k is the relative scaling of the content
    # block; it is not present in the idea's stated sink-only coordinate.
    k = np.log((1.0 - sa) / (1.0 - sb))
    corrected = m * dc2 + (m / D) * (ua - ub - g) ** 2
    stated = (m / D) * (dc2 + ds2)
    return dict(full=full, content=dc2, sink=ds2, corrected=corrected,
                stated=stated, content_logratio_mean=g, content_scale=k)


def math_sweep():
    invariant = []
    for m in [2, 3, 7, 15]:
        c1 = rng.dirichlet(np.ones(m) * .7)
        c2 = rng.dirichlet(np.ones(m) * .7)
        vals = [distances(row(s1, c1), row(s2, c2))['content']
                for s1, s2 in [(.1, .2), (.5, .8), (.9, .97)]]
        invariant.append({'m': m, 'predicted_range': 0.0,
                          'observed_range': float(max(vals) - min(vals))})

    sink_scaling = []
    for m in [2, 3, 7, 15, 31]:
        c = rng.dirichlet(np.ones(m) * .7)
        vals = []
        for _ in range(100):
            s1, s2 = rng.uniform(.05, .95, 2)
            d = distances(row(s1, c), row(s2, c))
            vals.append(d['full'] / d['sink'])
        sink_scaling.append({'m': m, 'predicted': 2*m/(m+1),
                             'observed_mean': float(np.mean(vals)),
                             'observed_sd': float(np.std(vals))})

    check = []
    for m in [2, 3, 7, 15, 31]:
        corrected_err, stated_err = [], []
        for _ in range(300):
            a = rng.dirichlet(np.ones(m+1) * .7)
            b = rng.dirichlet(np.ones(m+1) * .7)
            d = distances(a, b)
            corrected_err.append(abs(d['full'] - d['corrected']) / (d['full'] + EPS))
            stated_err.append(abs(d['full'] - d['stated']) / (d['full'] + EPS))
        check.append({'m': m, 'corrected_max_relative_error': float(max(corrected_err)),
                      'stated_median_relative_error': float(np.median(stated_err)),
                      'stated_mean_relative_error': float(np.mean(stated_err))})
    return {'content_sink_invariance': invariant,
            'equal_content_sink_scaling': sink_scaling,
            'decomposition_check': check}


def cosine_distance(a, b):
    return 1 - float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))


def js(a, b):
    q = .5 * (a + b)
    return float(.5*np.sum(a*np.log((a+EPS)/(q+EPS))) +
                 .5*np.sum(b*np.log((b+EPS)/(q+EPS))))


def pruning_toy():
    m = 5
    ca = np.array([.55, .25, .12, .06, .02])
    cb = np.array([.02, .06, .12, .25, .55])
    heads = [row(.95, ca), row(.95, cb), row(.70, ca),
             row(.95, ca*np.array([1.01,.99,1,1,1])), row(.50, cb)]
    names = ['reference', 'same_sink_diff_content', 'same_content_diff_sink',
             'duplicate', 'unrelated']
    pairs = []
    for i in range(len(heads)):
        for j in range(i+1, len(heads)):
            d = distances(heads[i], heads[j])
            pairs.append({'pair': [i,j], 'names': [names[i],names[j]],
                          'cosine': cosine_distance(heads[i],heads[j]),
                          'js': js(heads[i],heads[j]), 'content': d['content'],
                          'sink': d['sink'], 'full': d['full']})
    raw = sorted(pairs, key=lambda x: x['cosine'])
    admissible = [p for p in pairs if p['content'] < .02 and p['sink'] < .05]
    return {'pairs': pairs, 'raw_cosine_nearest': raw[:4],
            'channel_admissible': admissible,
            'thresholds': {'content': .02, 'sink': .05},
            'intended_redundant_pair': [0,3]}


def main():
    out = {'seed': SEED, 'math': math_sweep(), 'pruning_toy': pruning_toy()}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))


if __name__ == '__main__':
    main()