import json
import math
from pathlib import Path
import numpy as np


def paper_jacobian(t):
    t1, t2, t3, t4, t5, t6, t7, t8, t9 = np.asarray(t, dtype=float)
    return np.array([
        [-t1+t6+t7, 0, 0, 0, 0, t1-t6+t7, t1+t6-t7, 0, 0],
        [0, -t2+t3+t9, t2-t3+t9, 0, 0, 0, 0, 0, t2+t3-t9],
        [0, 0, 0, -t4+t5+t8, t4-t5+t8, 0, 0, t4+t5-t8, 0],
        [0, 0, 0, 0, 0, 0, -t7+t8+t9, t7-t8+t9, t7+t8-t9],
    ], dtype=float)


def numerical_rank(a, delta=1e-7):
    singular = np.linalg.svd(np.asarray(a), compute_uv=False)
    return int(np.sum(singular > delta)), singular


def triangle_area(lengths):
    a, b, c = map(float, lengths)
    s = (a + b + c) / 2.0
    return math.sqrt(max(s * (s-a) * (s-b) * (s-c), 1e-14))


def triangle_area_gradient(lengths):
    # Exact gradient of Heron's area with respect to (a,b,c).
    a, b, c = map(float, lengths)
    area = triangle_area(lengths)
    return np.array([
        a * (b*b + c*c - a*a),
        b * (a*a + c*c - b*b),
        c * (a*a + b*b - c*c),
    ]) / (8.0 * area)


def finite_difference_check():
    lengths = np.array([1.1, 1.3, 1.6])
    eps = 1e-6
    analytic = triangle_area_gradient(lengths)
    numeric = np.array([
        (triangle_area(lengths + eps*np.eye(3)[i]) -
         triangle_area(lengths - eps*np.eye(3)[i])) / (2*eps)
        for i in range(3)
    ])
    return {
        "analytic": analytic.tolist(), "finite_difference": numeric.tolist(),
        "max_abs_error": float(np.max(np.abs(analytic - numeric))),
    }


def make_candidates(seed=7, n_edges=30, n_candidates=180):
    rng = np.random.default_rng(seed)
    rows, lengths = [], []
    for _ in range(n_candidates):
        center = int(rng.integers(0, n_edges))
        tri = np.array([(center + int(rng.integers(-3, 4))) % n_edges,
                        int(rng.integers(0, n_edges)), int(rng.integers(0, n_edges))])
        while len(set(tri.tolist())) < 3:
            tri = rng.choice(n_edges, size=3, replace=False)
        pts = rng.normal(size=(3, 2))
        ls = np.array([np.linalg.norm(pts[1]-pts[2]),
                       np.linalg.norm(pts[0]-pts[2]),
                       np.linalg.norm(pts[0]-pts[1])])
        row = np.zeros(n_edges)
        row[tri] = triangle_area_gradient(ls)
        rows.append(row)
        lengths.append(ls)
    return np.asarray(rows), np.asarray(lengths)


def logdet_information(rows, lam=1e-3):
    gram = lam * np.eye(rows.shape[1]) + rows.T @ rows
    sign, value = np.linalg.slogdet(gram)
    return float(value) if sign > 0 else -np.inf


def greedy_select(rows, budget, lam=1e-3):
    selected, remaining = [], list(range(len(rows)))
    info = lam * np.eye(rows.shape[1])
    rank_trace = []
    for _ in range(min(budget, len(rows))):
        inv = np.linalg.inv(info)
        gains = [math.log1p(float(rows[i] @ inv @ rows[i])) for i in remaining]
        pos = int(np.argmax(gains))
        idx = remaining.pop(pos)
        selected.append(idx)
        info += np.outer(rows[idx], rows[idx])
        rank_trace.append(numerical_rank(rows[selected])[0])
    return selected, rank_trace


def random_select(rows, budget, rng):
    return rng.choice(len(rows), size=budget, replace=False).tolist()


def mean_cosine(rows):
    z = rows / np.maximum(np.linalg.norm(rows, axis=1, keepdims=True), 1e-12)
    vals = [abs(float(z[i] @ z[j])) for i in range(len(z)) for j in range(i)]
    return float(np.mean(vals)) if vals else 0.0


def metrics(rows, selected, lam=1e-3):
    chosen = rows[selected]
    rank, singular = numerical_rank(chosen)
    return {
        "rank": rank,
        "logdet_information": logdet_information(chosen, lam),
        "mean_pairwise_cosine_abs": mean_cosine(chosen),
        "smallest_singular_value": float(singular[-1]) if len(singular) else 0.0,
    }


def main():
    # Explicit witness from the paper: columns 1,2,4,7 are 1-indexed.
    witness = paper_jacobian([2, 1, 0, 1, 0, 0, 1, 0, 0])
    witness_rank, witness_singular = numerical_rank(witness)
    witness_minor_det = float(np.linalg.det(witness[:, [0, 1, 3, 6]]))

    rows, lengths = make_candidates()
    rng = np.random.default_rng(123)
    fractions = [0.10, 0.25, 0.50]
    sweep = {}
    for fraction in fractions:
        budget = int(round(fraction * len(rows)))
        selected, rank_trace = greedy_select(rows, budget)
        random_metrics = [metrics(rows, random_select(rows, budget, rng)) for _ in range(40)]
        keys = ["rank", "logdet_information", "mean_pairwise_cosine_abs", "smallest_singular_value"]
        sweep[str(int(100*fraction)) + "%"] = {
            "budget": budget,
            "greedy": metrics(rows, selected),
            "greedy_rank_trace": rank_trace,
            "random_mean": {k: float(np.mean([x[k] for x in random_metrics])) for k in keys},
            "random_std": {k: float(np.std([x[k] for x in random_metrics])) for k in keys},
            "message_fraction": budget / len(rows),
        }

    output = {
        "finite_difference_gradient_check": finite_difference_check(),
        "paper_witness": {
            "rank": witness_rank,
            "singular_values": witness_singular.tolist(),
            "selected_minor_determinant": witness_minor_det,
        },
        "candidate_pool": {"candidates": len(rows), "global_edges": rows.shape[1]},
        "all_candidates": metrics(rows, list(range(len(rows)))),
        "budget_sweep": sweep,
        "interpretation": "Greedy selects rank-increasing rows and maximizes regularized global edge information; random is the equal-budget control.",
    }
    Path("results.json").write_text(json.dumps(output, indent=2))
    print(json.dumps(output, indent=2))


if __name__ == "__main__":
    main()