import numpy as np
import pandas as pd
from collections import defaultdict

def compute_R_p(spatial_signature, p):
    n = spatial_signature.shape[0]
    angles = []
    for i in range(n):
        for j in range(i + 1, n):
            v = spatial_signature[i] - spatial_signature[j]
            base = np.arctan2(v[1], v[0])
            theta = (base + 0.5 * np.pi) % (2 * np.pi)
            angles.append(theta)
            angles.append((theta + np.pi) % (2 * np.pi))
    angles = np.sort(np.array(angles))

    diff = np.diff(np.concatenate([angles, angles[:1] + 2 * np.pi]))
    lambdas = diff
    K = len(lambdas)

    repeated_lambdas = np.tile(lambdas, 2)
    cum_lambdas = np.concatenate([[0], np.cumsum(repeated_lambdas)])

    total_area = 0.0
    for k in range(K):
        b = cum_lambdas[k + p] - cum_lambdas[k]
        a = cum_lambdas[k + K] - cum_lambdas[k + K - p]
        ell = lambdas[k]
        t_star = (b - a) / 2.0

        if t_star <= 0:
            area = b * ell - 0.5 * ell**2
        elif t_star >= ell:
            area = a * ell + 0.5 * ell**2
        else:
            area = (
                a * t_star
                + 0.5 * t_star**2
                + b * (ell - t_star)
                - 0.5 * (ell**2 - t_star**2)
            )
        total_area += area

    E_rho_p = total_area / (2 * np.pi)
    R_p = E_rho_p / (np.pi / 4)
    return R_p

def get_spatial_signature(values_dict, dataset, semivalue, i):
    u1 = np.array(values_dict[dataset][semivalue][i][1])
    u2 = np.array(values_dict[dataset][semivalue][i][2])
    return np.column_stack((u1, u2))

def summarize_R_p(values_dict, p):
    out = {}
    for dataset, semivalue_dict in values_dict.items():
        out[dataset] = {}
        for semivalue, replicates in semivalue_dict.items():
            n_reps = len(replicates)
            Rps = []
            for i in range(n_reps):
                sig = get_spatial_signature(values_dict, dataset, semivalue, i)
                Rps.append(compute_R_p(sig, p=p))
            Rps = np.array(Rps)
            mean = Rps.mean()
            se = Rps.std(ddof=1) / np.sqrt(n_reps)
            out[dataset][semivalue] = (mean, se)

    df = pd.DataFrame.from_dict(
        {(ds, sv): out[ds][sv] for ds in out for sv in out[ds]},
        orient="index",
        columns=["mean", "se"],
    )
    df.index = pd.MultiIndex.from_tuples(df.index, names=["Dataset", "Semivalue"])
    return df

def build_flattened_Rp_df(all_values, p_values):
    all_flat_data = defaultdict(dict)

    for p in p_values:
        df = summarize_R_p(all_values, p)
        for (dataset, method), row in df.iterrows():
            all_flat_data[dataset][f"R{p}_{method}_mean"] = row["mean"]
            all_flat_data[dataset][f"R{p}_{method}_se"] = row["se"]

    flat_df = pd.DataFrame.from_dict(all_flat_data, orient="index").reset_index()
    flat_df = flat_df.rename(columns={"index": "Dataset"})
    return flat_df
