import numpy as np
import json
import scipy.special
import string
from copy import deepcopy

# ── Helpers ───────────────────────────────────────────────────────────────────

def pae_to_logits(pae_matrix):
    """
    Approximate logits from a final PAE matrix using AF2 bin structure.
    Places all probability mass in the bin closest to each PAE value.
    
    NOTE: This is an approximation — real logits would have a spread of
    probability across bins. This assumes maximum confidence in the PAE value.
    """
    breaks = np.linspace(0, 31.75, 64)
    num_res = pae_matrix.shape[0]
    num_bins = len(breaks)

    bin_indices = np.argmin(
        np.abs(pae_matrix[:, :, None] - breaks[None, None, :]),
        axis=-1
    )

    logits = np.full((num_res, num_res, num_bins), -1e9)
    rows, cols = np.meshgrid(np.arange(num_res), np.arange(num_res), indexing='ij')
    logits[rows, cols, bin_indices] = 0.0

    return logits, breaks


def load_af3_result(af3_json_path):
    """
    Load an AF3 JSON output file and convert it into the result dictionary
    format expected by get_chain_and_interface_metrics.

    NOTE: Assumes af3_json['pae'] is a 2D list/array of shape [N, N].
    Check this matches your actual AF3 output structure.
    """
    with open(af3_json_path) as f:
        af3_json = json.load(f)

    pae = np.array(af3_json['pae'])
    logits, breaks = pae_to_logits(pae)

    result = {
        'predicted_aligned_error': {
            'logits': logits,
            'breaks': breaks
        },
        'contact_probs': np.array(af3_json['contact_probs']),
    }

    # NOTE: token_chain_ids is likely a list of chain labels e.g. ['A','A','B','B',...]
    # We need to convert this to integer asym_id e.g. [0,0,1,1,...]
    # Check this matches your actual AF3 output structure.
    chain_labels = af3_json['token_chain_ids']
    unique_chains = sorted(set(chain_labels), key=chain_labels.index)
    chain_to_int = {c: i for i, c in enumerate(unique_chains)}
    asym_id = np.array([chain_to_int[c] for c in chain_labels])

    return result, asym_id


def get_contact_map_af3(result, dist=8.0):
    """
    Replacement for get_contact_map using AF3 contact_probs directly.
    
    NOTE: The original function returns soft contact probabilities weighted
    by distance bins < 8A. Here we use contact_probs directly, which in AF3
    represents the probability of a contact — but the exact distance threshold
    AF3 uses internally may differ from 8A. Worth checking in the AF3 docs.
    """
    return np.array(result['contact_probs'])


# ── Core scoring functions (unchanged from original) ─────────────────────────

def _calculate_bin_centers(breaks):
    step = breaks[1] - breaks[0]
    bin_centers = breaks + step / 2
    return bin_centers


def predicted_tm_score_modified(logits, breaks, residue_weights=None,
                                 asym_id=None, pair_residue_weights=None):
    _np, _softmax = np, scipy.special.softmax

    if residue_weights is None:
        residue_weights = _np.ones(logits.shape[0])

    bin_centers = _calculate_bin_centers(breaks)
    num_res = residue_weights.shape[0]
    clipped_num_res = _np.maximum(num_res, 19)
    d0 = 1.24 * (clipped_num_res - 15) ** (1. / 3) - 1.8

    probs = _softmax(logits, axis=-1)
    tm_per_bin = 1. / (1 + _np.square(bin_centers) / _np.square(d0))
    predicted_tm_term = (probs * tm_per_bin).sum(-1)

    if asym_id is None:
        pair_mask = _np.full((num_res, num_res), True)
    else:
        pair_mask = asym_id[:, None] != asym_id[None, :]

    predicted_tm_term *= pair_mask

    if pair_residue_weights is None:
        pair_residue_weights = pair_mask * (
            residue_weights[None, :] * residue_weights[:, None]
        )
    normed_residue_mask = pair_residue_weights / (
        1e-8 + pair_residue_weights.sum(-1, keepdims=True)
    )

    per_alignment = (predicted_tm_term * normed_residue_mask).sum(-1)
    residuewise_iptm = per_alignment * residue_weights

    return residuewise_iptm


def get_chain_indices(asym_id):
    """Get start and end indices for each chain."""
    chain_ids = []
    seen = []
    for cid in asym_id:
        if cid not in seen:
            seen.append(cid)
            chain_ids.append(cid)

    indices = []
    for cid in chain_ids:
        positions = np.where(asym_id == cid)[0]
        indices.append((positions[0], positions[-1]))
    return indices


def get_actifptm_probs(result, asym_id, cmap, start_i, end_i, start_j, end_j):
    total_length = len(asym_id)
    outputs = deepcopy(result)

    cmap_copy = np.zeros((total_length, total_length))
    cmap_copy[start_i:end_i + 1, start_j:end_j + 1] = cmap[
        start_i:end_i + 1, start_j:end_j + 1
    ]
    cmap_copy[start_j:end_j + 1, start_i:end_i + 1] = cmap[
        start_j:end_j + 1, start_i:end_i + 1
    ]

    if end_i == end_j == total_length - 1 and start_i == start_j == 0:
        pair_mask = asym_id[:, None] != asym_id[None, :]
        cmap_copy *= pair_mask

    seq_mask = np.full(total_length, 0, dtype=float)
    seq_mask[np.concatenate((np.arange(start_i, end_i + 1),
                              np.arange(start_j, end_j + 1)))] = 1

    pae = {"residue_weights": seq_mask,
           **outputs["predicted_aligned_error"],
           "asym_id": asym_id}
    residuewise_actifptm = predicted_tm_score_modified(**pae,
                                                        pair_residue_weights=cmap_copy)
    return residuewise_actifptm


def get_pairwise_iptm(result, asym_id, start_i, end_i, start_j, end_j):
    total_length = len(asym_id)
    outputs = deepcopy(result)

    seq_mask = np.full(total_length, 0, dtype=float)
    seq_mask[np.concatenate((np.arange(start_i, end_i + 1),
                              np.arange(start_j, end_j + 1)))] = 1

    pae = {"residue_weights": seq_mask,
           **outputs["predicted_aligned_error"],
           "asym_id": asym_id}
    return predicted_tm_score_modified(**pae)


def get_per_chain_ptm(result, cmap, start_i, end_i):
    outputs = deepcopy(result)
    total_length = outputs["predicted_aligned_error"]["logits"].shape[0]

    seq_mask = np.full(total_length, 0, dtype=float)
    seq_mask[start_i:end_i + 1] = 1

    cmap_chain = np.zeros((total_length, total_length))
    cmap_chain[start_i:end_i + 1, start_i:end_i + 1] = cmap[
        start_i:end_i + 1, start_i:end_i + 1
    ]

    pae = {"residue_weights": seq_mask,
           **outputs["predicted_aligned_error"],
           "asym_id": None}
    return round(float(predicted_tm_score_modified(**pae,
                                                    pair_residue_weights=cmap_chain
                                                    ).max()), 3)


# ── Simplified main function ──────────────────────────────────────────────────

def get_chain_and_interface_metrics(result, asym_id):
    """
    Calculates interface and interchain PTM scores for each chain pair,
    using contact probabilities.

    Args:
        result: Result dictionary constructed by load_af3_result()
        asym_id: Integer array of shape [N] indicating chain membership,
                 constructed by load_af3_result()
    Returns:
        Dictionary with pairwise_actifptm, pairwise_iptm, per_chain_ptm,
        and overall actifptm.
    """
    if len(asym_id.shape) > 1:
        asym_id = asym_id[0]

    full_length = len(asym_id)
    output = {'pairwise_actifptm': {}, 'pairwise_iptm': {}, 'per_chain_ptm': {}}
    chain_starts_ends = get_chain_indices(asym_id)
    chain_labels = list(string.ascii_uppercase)

    cmap = get_contact_map_af3(result)
    results = {'predicted_aligned_error': deepcopy(result['predicted_aligned_error'])}

    for i, (start_i, end_i) in enumerate(chain_starts_ends):
        chain_label_i = chain_labels[i % len(chain_labels)]
        for j, (start_j, end_j) in enumerate(chain_starts_ends):
            chain_label_j = chain_labels[j % len(chain_labels)]
            if i < j:
                key = f"{chain_label_i}-{chain_label_j}"

                residuewise_actifptm = get_actifptm_probs(results, asym_id, cmap,
                                                           start_i, end_i,
                                                           start_j, end_j)
                output['pairwise_actifptm'][key] = round(
                    float(residuewise_actifptm.max()), 3
                )

                residuewise_iptm = get_pairwise_iptm(results, asym_id,
                                                      start_i, end_i,
                                                      start_j, end_j)
                output['pairwise_iptm'][key] = round(float(residuewise_iptm.max()), 3)

        output['per_chain_ptm'][chain_label_i] = get_per_chain_ptm(results, cmap,
                                                                     start_i, end_i)

    output['actifptm'] = round(float(get_actifptm_probs(results, asym_id, cmap,
                                                          0, full_length - 1,
                                                          0, full_length - 1).max()), 3)
    return output


def load_colabfold_result(cf_json_path):
    with open(cf_json_path) as f:
        cf_json = json.load(f)

    pae = np.array(cf_json['pae'])
    total_residues = pae.shape[0]
    num_chains = len(cf_json['per_chain_ptm'])
    
    # Infer chain length from PAE shape
    # NOTE: this assumes all chains are equal length — if not, 
    # you'll need to supply chain_lengths manually
    assert total_residues % num_chains == 0, \
        f"Total residues {total_residues} not evenly divisible by {num_chains} chains"
    chain_length = total_residues // num_chains
    chain_lengths = [chain_length] * num_chains
    print(f"Inferred: {num_chains} chains x {chain_length} residues = {total_residues} total")

    logits, breaks = pae_to_logits(pae)
    result = {
        'predicted_aligned_error': {
            'logits': logits,
            'breaks': breaks
        },
        'contact_probs': (pae < 6.0).astype(float),
    }

    asym_id = np.concatenate([
        np.full(length, i) for i, length in enumerate(chain_lengths)
    ])

    return result, asym_id, cf_json


def compare_scores(computed, precomputed):
    """
    Compare our approximated scores against the pre-computed ones
    already in the ColabFold JSON.
    """
    print("\n=== actifPTM comparison ===")
    print(f"{'Metric':<20} {'Pre-computed':>15} {'Ours':>15} {'Diff':>10}")
    print("-" * 62)

    # Overall actifptm
    pre = precomputed.get('actifptm')
    ours = computed.get('actifptm')
    if pre is not None and ours is not None:
        print(f"{'actifptm':<20} {pre:>15.3f} {ours:>15.3f} {abs(pre-ours):>10.3f}")

    # Pairwise actifptm
    print("\n--- Pairwise actifPTM ---")
    print(f"{'Pair':<20} {'Pre-computed':>15} {'Ours':>15} {'Diff':>10}")
    print("-" * 62)
    for key, pre in precomputed.get('pairwise_actifptm', {}).items():
        ours = computed.get('pairwise_actifptm', {}).get(key)
        if ours is not None:
            print(f"{key:<20} {pre:>15.3f} {ours:>15.3f} {abs(pre-ours):>10.3f}")

    # Pairwise iptm
    print("\n--- Pairwise iPTM ---")
    print(f"{'Pair':<20} {'Pre-computed':>15} {'Ours':>15} {'Diff':>10}")
    print("-" * 62)
    for key, pre in precomputed.get('pairwise_iptm', {}).items():
        ours = computed.get('pairwise_iptm', {}).get(key)
        if ours is not None:
            print(f"{key:<20} {pre:>15.3f} {ours:>15.3f} {abs(pre-ours):>10.3f}")


if __name__ == "__main__":
    import sys

    mode = sys.argv[1]  # "af3" or "colabfold"

    if mode == "af3":
        json_path = sys.argv[2]
        result, asym_id = load_af3_result(json_path)
        metrics = get_chain_and_interface_metrics(result, asym_id)
        print(json.dumps(metrics, indent=2))

    elif mode == "colabfold":
        json_path = sys.argv[2]
        result, asym_id, cf_json = load_colabfold_result(json_path)
        metrics = get_chain_and_interface_metrics(result, asym_id)
        compare_scores(metrics, cf_json)