"""
Sim-CoT Information Load Simulator

Since Sim-CoT uses rigid one-to-one segmentation (evenly dividing N_c CoT
tokens into N_l segments, each mapped to one latent token), we can simulate
its information load distribution without an actual Sim-CoT checkpoint.

This script reads CoT lengths from PMPS results and computes what the
Sim-CoT distribution would look like.

Usage:
    python analysis/simulate_simcot_distribution.py \
        --pmps_dir analysis/results/pmps \
        --num_latent 6 \
        --output_dir analysis/results/simcot_simulated
"""

import os
import json
import argparse
import numpy as np


def simulate_simcot_info_load(cot_lengths, num_latent):
    """
    Simulate Sim-CoT's rigid segmentation information load.

    Sim-CoT divides the CoT into num_latent equal segments.
    Each latent token gets floor(N_c / N_l) tokens, with the last
    latent getting the remainder.

    However, for samples where N_c < N_l, some latents get 0 tokens
    (the core information bottleneck problem).

    Returns:
        info_loads: list of (num_latent,) arrays
    """
    info_loads = []
    for N_c in cot_lengths:
        load = np.zeros(num_latent)
        if N_c <= 0:
            info_loads.append(load.tolist())
            continue

        if N_c >= num_latent:
            # Normal case: divide evenly
            base_size = N_c // num_latent
            remainder = N_c % num_latent
            for i in range(num_latent):
                # First `remainder` latents get one extra token
                load[i] = base_size + (1 if i < remainder else 0)
        else:
            # N_c < N_l: some latents get nothing (surplus latents)
            for i in range(N_c):
                load[i] = 1
            # Remaining latents get 0 → unsupervised

        info_loads.append(load.tolist())

    return info_loads


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--pmps_dir', type=str, required=True,
                        help='PMPS results directory (to read CoT lengths)')
    parser.add_argument('--num_latent', type=int, default=6)
    parser.add_argument('--output_dir', type=str, default='analysis/results/simcot_simulated')
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    # Read CoT lengths from PMPS results
    cot_lengths_path = os.path.join(args.pmps_dir, 'cot_lengths.json')
    if os.path.exists(cot_lengths_path):
        with open(cot_lengths_path, 'r') as f:
            cot_lengths = json.load(f)
    else:
        # Fallback: try to read from distribution metrics
        dist_path = os.path.join(args.pmps_dir, 'distribution_metrics.json')
        with open(dist_path, 'r') as f:
            dist_metrics = json.load(f)
        # Estimate CoT length from info load sums
        cot_lengths = [int(sum(m['info_load_per_latent'])) for m in dist_metrics]

    print(f'Loaded {len(cot_lengths)} CoT lengths')
    print(f'  Mean CoT length: {np.mean(cot_lengths):.1f}')
    print(f'  Min/Max: {np.min(cot_lengths)}/{np.max(cot_lengths)}')

    # Simulate Sim-CoT info loads
    info_loads = simulate_simcot_info_load(cot_lengths, args.num_latent)

    # Compute distribution metrics
    dist_metrics = []
    for i, (loads, cot_len) in enumerate(zip(info_loads, cot_lengths)):
        loads_arr = np.array(loads)
        cv = float(np.std(loads_arr) / (np.mean(loads_arr) + 1e-8))
        dist_metrics.append({
            'sample_idx': i,
            'info_load_per_latent': loads,
            'info_load_cv': cv,
            'cot_length': cot_len,
            # Sim-CoT has no prototype concept
            'proto_usage_latent': [1.0 / args.num_latent] * args.num_latent,
            'proto_entropy_latent': float(np.log(args.num_latent)),
            'proto_usage_cot': [1.0 / args.num_latent] * args.num_latent,
            'proto_entropy_cot': float(np.log(args.num_latent)),
            'max_entropy': float(np.log(args.num_latent)),
            'avg_match_entropy': 0.0,  # rigid 1-to-1 → entropy = 0
            'match_entropies': [0.0] * cot_len,
        })

    # Also create a fake metrics.json for the collapse analysis pipeline
    # (we need cos_sim_matrix and effective_rank — use zeros as placeholders
    #  unless real Sim-CoT checkpoint data is available)
    collapse_metrics = []
    for i in range(len(info_loads)):
        N_l = args.num_latent
        collapse_metrics.append({
            'sample_idx': i,
            'cos_sim_matrix': np.eye(N_l).tolist(),  # placeholder
            'avg_off_diag_sim': 0.0,
            'std_off_diag_sim': 0.0,
            'singular_values': [1.0] * N_l,
            'singular_values_normalized': [1.0 / N_l] * N_l,
            'effective_rank': float(N_l),
            'explained_variance': [1.0 / N_l] * N_l,
            'cumulative_variance': [i / N_l for i in range(1, N_l + 1)],
        })

    # Save
    with open(os.path.join(args.output_dir, 'distribution_metrics.json'), 'w') as f:
        json.dump(dist_metrics, f, indent=2)

    with open(os.path.join(args.output_dir, 'metrics.json'), 'w') as f:
        json.dump(collapse_metrics, f, indent=2)

    # Save simulated latent hiddens (zeros placeholder)
    D = 768  # placeholder hidden dim
    fake_hiddens = np.zeros((len(info_loads), args.num_latent, D))
    np.save(os.path.join(args.output_dir, 'latent_hiddens.npy'), fake_hiddens)

    # Summary
    all_cvs = [m['info_load_cv'] for m in dist_metrics]
    print(f'\nSimulated Sim-CoT distribution:')
    print(f'  Mean Info Load CV: {np.mean(all_cvs):.4f} (higher = less uniform)')
    print(f'  Samples with surplus latents (N_c < N_l): '
          f'{sum(1 for l in cot_lengths if l < args.num_latent)}')
    print(f'\nSaved to: {args.output_dir}')


if __name__ == '__main__':
    main()