import numpy as np
import os
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from scipy.special import logsumexp
from parismc import Sampler, SamplerConfig
import copy

# ==============================================================================
# 1. Problem Definition
# ==============================================================================

MODES = np.array([
    [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
])
WEIGHTS = np.array([1.0] * len(MODES))
LOG_WEIGHTS = np.log(WEIGHTS)

def log_density(final_states_raw):
    final_states_in = (final_states_raw * 1.4) - 0.2
    log_likelihoods = []
    for mode in MODES:
        distance = np.linalg.norm(final_states_in - mode, axis=1)
        distance = np.nan_to_num(distance, nan=1e6, posinf=1e6, neginf=1e6)
        log_prob = -400 * distance**2
        log_likelihoods.append(log_prob)
    
    log_likelihoods = np.array(log_likelihoods)
    log_likelihoods = LOG_WEIGHTS[:, None] + log_likelihoods
    total_log_likelihood = logsumexp(log_likelihoods, axis=0)
    if total_log_likelihood.shape[0] == 1:
        return total_log_likelihood[0]
    return total_log_likelihood

def prior_transform(u):
    return u

# ==============================================================================
# 2. Hook Class
# ==============================================================================

class TrajectoryHook:
    def __init__(self, n_seed, initial_points):
        self.n_seed = n_seed
        self.all_samples = [] # (iter, x, y, proc_id)
        self.merge_events = [] # (iter, x, y)
        
        # Track state
        self.prev_n_proc = n_seed
        # We need to track the "last known location" of each process to identify deaths
        # Initial means are the initial points
        self.prev_means = [initial_points[i] for i in range(n_seed)]
        
        # Track element counts to know what is new
        self.last_counts = {i: 0 for i in range(n_seed)}

    def __call__(self, sampler, iter_idx):
        # 1. Collect NEW samples
        # Iterate over active processes (0 to sampler.n_proc - 1)
        for p_idx in range(sampler.n_proc):
            curr_count = sampler.element_num_list[p_idx]
            prev_cnt = self.last_counts.get(p_idx, 0)
            
            if curr_count > prev_cnt:
                # Get new points
                new_pts = sampler.searched_points_list[p_idx][prev_cnt:curr_count]
                
                # Subsample if too many? (e.g. take every 10th if iter is large)
                # For 5000 iters * 1 point/iter, it's fine.
                for pt in new_pts:
                    self.all_samples.append((iter_idx, pt[0], pt[1], p_idx))
            
            self.last_counts[p_idx] = curr_count
            
        # 2. Detect Merge (Did n_proc drop compared to PREVIOUS call?)
        # Note: callback is called at START of iter i.
        # So sampler.n_proc reflects the state after iter i-1.
        
        current_n_proc = sampler.n_proc
        if current_n_proc < self.prev_n_proc:
            print(f"Merge detected at start of iter {iter_idx}. Active: {self.prev_n_proc} -> {current_n_proc}")
            
            # Identify who died using PREVIOUS means vs CURRENT means
            current_means = sampler.now_means
            
            # Simple matching: find which of `self.prev_means` is furthest from any `current_mean`
            # This assumes the survivor didn't jump too far.
            
            if len(current_means) > 0:
                matched_indices = []
                for cm in current_means:
                    dists = [np.linalg.norm(cm - bm) for bm in self.prev_means]
                    best_idx = np.argmin(dists)
                    matched_indices.append(best_idx)
                
                for idx in range(self.prev_n_proc):
                    if idx not in matched_indices:
                        dead_loc = self.prev_means[idx]
                        self.merge_events.append((iter_idx, dead_loc[0], dead_loc[1]))
            else:
                # If no current means (rare/impossible if n_proc > 0), just log something
                pass

        # 3. Update State for NEXT call
        self.prev_n_proc = current_n_proc
        # Snapshot current means
        if hasattr(sampler, 'now_means') and len(sampler.now_means) > 0:
            self.prev_means = [m.copy() for m in sampler.now_means]
        # Else keep prev_means (e.g. at iter 0 if now_means not populated yet)

# ==============================================================================
# 3. Experiment Runner
# ==============================================================================

def run_callback_experiment(merge_type, run_id, total_iters=5000):
    ndim = 10
    n_seed = 2
    sigma = 0.01
    savepath = f'./vis_results_cb/run_{run_id}_{merge_type}/'
    os.makedirs(savepath, exist_ok=True)
    
    init_cov_list = [sigma**2 * np.eye(ndim) for _ in range(n_seed)]
    
    # Config: same as merge_exp.py
    config = SamplerConfig(
        merge_confidence=0.9,
        alpha=1000,
        trail_size=int(1e3),
        boundary_limiting=True,
        use_beta=True,
        integral_num=int(1e5),
        gamma=100,
        exclude_scale_z=np.inf,
        use_pool=False,
        n_pool=4,
        merge_type=merge_type,
        stop_on_merge=False # Continue!
    )
    
    sampler = Sampler(ndim, n_seed, log_density, init_cov_list, prior_transform, config)
    
    external_lhs_points = np.zeros((n_seed, ndim)) + 0.4
    external_lhs_points[1, :] = 0.6
    external_lhs_log_densities = log_density(external_lhs_points)
    
    # Create Hook
    hook = TrajectoryHook(n_seed, external_lhs_points)
    
    try:
        # Run ALL AT ONCE with callback
        sampler.run_sampling(
            num_iterations=total_iters,
            savepath=savepath,
            print_iter=total_iters+100, # Quiet
            stop_dlogZ=0.001,
            external_lhs_points=external_lhs_points,
            external_lhs_log_densities=external_lhs_log_densities,
            callback=hook
        )
    except Exception as e:
        print(f"Run {merge_type} stopped early: {e}")
        
    return np.array(hook.all_samples), hook.merge_events, external_lhs_points

# ==============================================================================
# 4. Main
# ==============================================================================

def main():
    merge_types = ['distance', 'single', 'multiple']
    
    fig, axes = plt.subplots(3, 1, figsize=(8, 18))
    target_mode = [0.5, 0.5]
    target_sigma = 0.02525
    
    for i, m_type in enumerate(merge_types):
        print(f"Running visualization for merge_type: {m_type}")
        samples, merges, init_pts = run_callback_experiment(m_type, i, total_iters=5000)
        
        ax = axes[i]
        
        # 0. Initial Points
        ax.scatter(init_pts[:, 0], init_pts[:, 1], marker='s', s=100, color='black', 
                   edgecolors='white', zorder=15, label='Start Points')
        
        # 1. Trajectory
        if len(samples) > 0:
            sc = ax.scatter(samples[:, 1], samples[:, 2], c=samples[:, 0], 
                            cmap='viridis', s=10, alpha=0.6, edgecolors='none')
            cbar = plt.colorbar(sc, ax=ax)
            cbar.set_label('Iteration')
            
        # 2. Target
        circle = Circle(target_mode, target_sigma, color='red', fill=False, 
                        linewidth=2, linestyle='--', label='Target 1-$\sigma$')
        ax.add_patch(circle)
        
        # 3. Merges
        for me in merges:
            ax.scatter(me[1], me[2], marker='*', s=300, color='crimson', 
                       edgecolors='black', zorder=10, label='Merge Event')
            ax.text(me[1]+0.02, me[2]+0.02, f"Iter {int(me[0])}", fontsize=10, color='crimson', weight='bold')
            
        ax.set_title(f"Merge Strategy: {m_type}")
        ax.set_xlim(0.2, 0.8)
        ax.set_ylim(0.2, 0.8)
        
        handles, labels = ax.get_legend_handles_labels()
        by_label = dict(zip(labels, handles))
        ax.legend(by_label.values(), by_label.keys(), loc='upper right')
        
        print(f"Finished {m_type}. Samples: {len(samples)}. Merges: {len(merges)}")

    plt.tight_layout()
    plt.savefig('visualize_merge_callback.png', dpi=300)
    print("Plot saved to visualize_merge_callback.png")

if __name__ == "__main__":
    main()
