"""
Simulation functions for BioGSP package.

Generate various spatial patterns for testing and demonstration.
"""

import numpy as np
import pandas as pd
from typing import Dict, List, Optional, Tuple, Union
from tqdm import tqdm


def simulate_multiscale(
    grid_size: int = 60,
    Ra_seq: Optional[np.ndarray] = None,
    n_steps: int = 10,
    n_centers: int = 1,
    outer_start: float = 40,
    outer_end_factor: float = 1.2,
    seed: Optional[int] = None,
    verbose: bool = True,
) -> Dict[str, pd.DataFrame]:
    """
    Simulate Multi-center Multi-scale Concentric Ring Patterns.
    
    Generate multi-center, multi-scale concentric ring simulation data.
    Creates patterns with inner circles and outer rings.
    
    Parameters
    ----------
    grid_size : int
        Size of the spatial grid.
    Ra_seq : np.ndarray, optional
        Vector of inner circle radii. Default: np.arange(2.5, 22.5, 2.5)
    n_steps : int
        Number of outer radius shrinkage steps.
    n_centers : int
        Number of circle centers.
    outer_start : float
        Fixed starting outer radius.
    outer_end_factor : float
        Outer radius shrinks to this factor * Ra.
    seed : int, optional
        Random seed for reproducible center placement.
    verbose : bool
        Whether to show progress.
        
    Returns
    -------
    dict
        Dictionary of DataFrames with X, Y coordinates and signal_1, signal_2.
        
    Examples
    --------
    >>> patterns = simulate_multiscale()
    >>> patterns = simulate_multiscale(grid_size=80, n_steps=8, n_centers=2)
    """
    if Ra_seq is None:
        Ra_seq = np.arange(2.5, 22.5, 2.5)
    Ra_seq = np.asarray(Ra_seq)
    
    if seed is not None:
        np.random.seed(seed)
    
    # Create grid
    X, Y = np.meshgrid(np.arange(1, grid_size + 1), np.arange(1, grid_size + 1))
    grid_X = X.flatten()
    grid_Y = Y.flatten()
    
    # Define circle centers
    if n_centers == 1:
        centers = [(grid_size / 2, grid_size / 2)]
    else:
        centers = [(np.random.randint(10, grid_size - 10), 
                   np.random.randint(10, grid_size - 10)) for _ in range(n_centers)]
    
    sim_list = {}
    total_iter = len(Ra_seq) * n_steps
    
    if verbose:
        print(f"Generating {total_iter} patterns...")
        pbar = tqdm(total=total_iter)
    
    for Ra in Ra_seq:
        # Outer radius shrinks from outer_start to Ra * outer_end_factor
        Rb_seq = np.linspace(outer_start, Ra * outer_end_factor, n_steps)
        
        for step_i, Rb in enumerate(Rb_seq, 1):
            circleA = np.zeros(len(grid_X), dtype=bool)
            circleB = np.zeros(len(grid_X), dtype=bool)
            
            for cx, cy in centers:
                dists = np.sqrt((grid_X - cx) ** 2 + (grid_Y - cy) ** 2)
                circleA = circleA | (dists <= Ra)
                circleB = circleB | ((dists > Ra) & (dists <= Rb))
            
            df = pd.DataFrame({
                "X": grid_X,
                "Y": grid_Y,
                "signal_1": circleA.astype(int),
                "signal_2": (circleB & ~circleA).astype(int),
            })
            
            sim_list[f"Ra_{Ra}_Step_{step_i}"] = df
            
            if verbose:
                pbar.update(1)
    
    if verbose:
        pbar.close()
        print("Multicenter multiscale simulation completed!")
    
    return sim_list


def simulate_stripe_patterns(
    grid_size: int = 100,
    gap_seq: Optional[List[float]] = None,
    width_seq: Optional[List[float]] = None,
    theta_seq: Optional[List[float]] = None,
    eps: float = 1e-9,
    verbose: bool = True,
) -> Dict[str, pd.DataFrame]:
    """
    Simulate Stripe Patterns.
    
    Generate stripe patterns with two parallel stripes separated by a gap.
    
    Parameters
    ----------
    grid_size : int
        Size of the spatial grid.
    gap_seq : list of float, optional
        Gap distances between stripe centers. Default: [10]
    width_seq : list of float, optional
        Stripe widths. Default: [5]
    theta_seq : list of float, optional
        Rotation angles in degrees. Default: [0]
    eps : float
        Small value for open boundary conditions.
    verbose : bool
        Whether to show progress.
        
    Returns
    -------
    dict
        Dictionary of DataFrames with patterns.
        
    Examples
    --------
    >>> patterns = simulate_stripe_patterns()
    >>> patterns = simulate_stripe_patterns(gap_seq=[10, 20], width_seq=[5, 10])
    """
    if gap_seq is None:
        gap_seq = [10]
    if width_seq is None:
        width_seq = [5]
    if theta_seq is None:
        theta_seq = [0]
    
    # Create grid
    X, Y = np.meshgrid(np.arange(1, grid_size + 1), np.arange(1, grid_size + 1))
    grid_X = X.flatten()
    grid_Y = Y.flatten()
    
    sim_list = {}
    total_iterations = len(gap_seq) * len(width_seq) * len(theta_seq)
    
    if verbose:
        print(f"Generating {total_iterations} stripe patterns...")
    
    iteration = 0
    
    def rotate(x, y, theta, cx, cy):
        """Rotate coordinates around center."""
        th = theta * np.pi / 180
        x_shift = x - cx
        y_shift = y - cy
        x_rot = np.cos(th) * x_shift + np.sin(th) * y_shift
        y_rot = -np.sin(th) * x_shift + np.cos(th) * y_shift
        return x_rot, y_rot
    
    for gap in gap_seq:
        for width in width_seq:
            for theta in theta_seq:
                x_left_line = grid_size / 2 - gap / 2
                x_right_line = grid_size / 2 + gap / 2
                y_center = grid_size / 2
                
                if gap == 0:
                    # Both stripes meet in the middle
                    x_rot, y_rot = rotate(grid_X, grid_Y, theta, grid_size / 2, y_center)
                    signal_1 = ((x_rot >= -width) & (x_rot <= 0)).astype(int)
                    signal_2 = ((x_rot > 0) & (x_rot <= width)).astype(int)
                else:
                    # Normal case: two separated stripes
                    x_rot_left, _ = rotate(grid_X, grid_Y, theta, x_left_line, y_center)
                    x_rot_right, _ = rotate(grid_X, grid_Y, theta, x_right_line, y_center)
                    
                    signal_1 = ((x_rot_left <= 0) & (x_rot_left >= -width)).astype(int)
                    signal_2 = ((x_rot_right >= 0) & (x_rot_right <= width)).astype(int)
                
                df = pd.DataFrame({
                    "X": grid_X,
                    "Y": grid_Y,
                    "signal_1": signal_1,
                    "signal_2": signal_2,
                    "gap_param": gap,
                    "width_param": width,
                    "theta_param": theta,
                })
                
                name = f"gap_{gap}_w_{width}_th_{theta}"
                sim_list[name] = df
                
                iteration += 1
                if verbose and iteration % max(1, total_iterations // 10) == 0:
                    print(f"Completed {iteration} of {total_iterations} patterns")
    
    if verbose:
        print("Stripe pattern simulation completed!")
    
    return sim_list


def simulate_multiscale_overlap(
    grid_size: int = 60,
    n_centers: int = 3,
    Ra_seq: Optional[List[float]] = None,
    Rb_seq: Optional[List[float]] = None,
    seed: Optional[int] = None,
    verbose: bool = True,
) -> Dict[str, pd.DataFrame]:
    """
    Simulate Multiple Center Patterns with Fixed Centers.
    
    Generate spatial patterns with multiple circular centers at fixed positions.
    
    Parameters
    ----------
    grid_size : int
        Size of the spatial grid.
    n_centers : int
        Number of pattern centers.
    Ra_seq : list of float, optional
        Inner circle radii. Default: [10, 5, 1]
    Rb_seq : list of float, optional
        Outer ring radii. Default: [10, 5, 1]
    seed : int, optional
        Random seed.
    verbose : bool
        Whether to show progress.
        
    Returns
    -------
    dict
        Dictionary of DataFrames with patterns.
    """
    if Ra_seq is None:
        Ra_seq = [10, 5, 1]
    if Rb_seq is None:
        Rb_seq = [10, 5, 1]
    
    # Create grid
    X, Y = np.meshgrid(np.arange(1, grid_size + 1), np.arange(1, grid_size + 1))
    grid_X = X.flatten()
    grid_Y = Y.flatten()
    
    if seed is not None:
        np.random.seed(seed)
    
    # Define centers
    if n_centers == 1:
        centers = [(grid_size / 2, grid_size / 2)]
    else:
        centers = [(np.random.randint(10, grid_size - 10),
                   np.random.randint(10, grid_size - 10)) for _ in range(n_centers)]
    
    sim_list = {}
    total_iterations = len(Ra_seq) * len(Rb_seq)
    
    if verbose:
        print(f"Generating {total_iterations} multiscale patterns...")
        pbar = tqdm(total=total_iterations)
    
    for Ra in Ra_seq:
        for Rb in Rb_seq:
            circleA = np.zeros(len(grid_X), dtype=bool)
            circleB = np.zeros(len(grid_X), dtype=bool)
            
            for cx, cy in centers:
                dists = np.sqrt((grid_X - cx) ** 2 + (grid_Y - cy) ** 2)
                circleA = circleA | (dists <= Ra)
                circleB = circleB | ((dists > Ra) & (dists <= Ra + Rb))
            
            signal_1 = circleA.astype(int)
            signal_2 = (circleB & ~circleA).astype(int)
            
            df = pd.DataFrame({
                "X": grid_X,
                "Y": grid_Y,
                "signal_1": signal_1,
                "signal_2": signal_2,
            })
            
            name = f"simulated_Ra_{Ra}_Rb_{Rb}"
            sim_list[name] = df
            
            if verbose:
                pbar.update(1)
    
    if verbose:
        pbar.close()
        print("Multiscale (fixed-center) simulation completed!")
    
    return sim_list


def simulate_moving_circles(
    grid_size: int = 60,
    radius_seq: Optional[np.ndarray] = None,
    n_steps: int = 10,
    center_distance: float = 30,
    radius2_factor: float = 1.5,
    seed: Optional[int] = None,
    verbose: bool = True,
) -> Dict[str, pd.DataFrame]:
    """
    Simulate Moving Circles Pattern.
    
    Generate patterns of two circles moving toward each other horizontally.
    
    Parameters
    ----------
    grid_size : int
        Size of the spatial grid.
    radius_seq : np.ndarray, optional
        Radii for circle 1. Default: np.arange(6, 15)
    n_steps : int
        Number of movement steps.
    center_distance : float
        Initial horizontal distance from midline.
    radius2_factor : float
        Circle 2 radius = radius_seq * radius2_factor.
    seed : int, optional
        Random seed.
    verbose : bool
        Whether to show progress.
        
    Returns
    -------
    dict
        Dictionary of DataFrames with patterns.
    """
    if radius_seq is None:
        radius_seq = np.arange(6, 15)
    radius_seq = np.asarray(radius_seq)
    
    if seed is not None:
        np.random.seed(seed)
    
    # Create grid
    X, Y = np.meshgrid(np.arange(1, grid_size + 1), np.arange(1, grid_size + 1))
    grid_X = X.flatten()
    grid_Y = Y.flatten()
    
    # Movement: two circles move toward the midline
    circle1_x = center_distance - np.linspace(center_distance / 2, 0, n_steps)
    circle2_x = center_distance + np.linspace(center_distance / 2, 0, n_steps)
    circle_y = grid_size / 2
    
    sim_list = {}
    total_iter = len(radius_seq) * n_steps
    
    if verbose:
        print(f"Generating {total_iter} patterns...")
        pbar = tqdm(total=total_iter)
    
    for r1 in radius_seq:
        r2 = r1 * radius2_factor
        
        for k in range(n_steps):
            cx1 = circle1_x[k]
            cx2 = circle2_x[k]
            cy = circle_y
            
            d1 = np.sqrt((grid_X - cx1) ** 2 + (grid_Y - cy) ** 2)
            d2 = np.sqrt((grid_X - cx2) ** 2 + (grid_Y - cy) ** 2)
            
            inside1 = d1 <= r1
            inside2 = d2 <= r2
            
            # Mutually exclusive: if overlap, assign to circle 1
            signal_1 = inside1.astype(int)
            signal_2 = (inside2 & ~inside1).astype(int)
            
            df = pd.DataFrame({
                "X": grid_X,
                "Y": grid_Y,
                "signal_1": signal_1,
                "signal_2": signal_2,
                "Ra": r1,
                "Step": k + 1,
            })
            
            sim_list[f"Ra_{r1}_Step_{k + 1}"] = df
            
            if verbose:
                pbar.update(1)
    
    if verbose:
        pbar.close()
        print("Moving-circles simulation completed!")
    
    return sim_list


def simulate_checkerboard(
    grid_size: int = 8,
    tile_size: int = 1,
) -> pd.DataFrame:
    """
    Simulate checkerboard pattern.
    
    Generate a checkerboard pattern with alternating signals.
    
    Parameters
    ----------
    grid_size : int
        Number of tiles per row/column.
    tile_size : int
        Resolution of each tile in pixels per side.
        
    Returns
    -------
    pd.DataFrame
        DataFrame with X, Y coordinates and signal_1, signal_2 patterns.
        
    Examples
    --------
    >>> df = simulate_checkerboard(grid_size=8, tile_size=10)
    """
    # Generate lattice grid
    xs = np.arange(1, grid_size * tile_size + 1)
    ys = np.arange(1, grid_size * tile_size + 1)
    X, Y = np.meshgrid(xs, ys)
    grid_X = X.flatten()
    grid_Y = Y.flatten()
    
    # Determine tile index for each coordinate
    tile_x = (grid_X - 1) // tile_size
    tile_y = (grid_Y - 1) // tile_size
    
    # Checkerboard pattern
    signal_1 = ((tile_x + tile_y) % 2 == 0).astype(int)
    signal_2 = ((tile_x + tile_y) % 2 == 1).astype(int)
    
    return pd.DataFrame({
        "X": grid_X,
        "Y": grid_Y,
        "signal_1": signal_1,
        "signal_2": signal_2,
    })
