"""
Kalman filter–based burst detection for (multi-channel) photon count data.
Includes loading from TTTR files using tttrlib and constructing binned time series.

Author: ChatGPT
License: MIT
"""

from __future__ import annotations
import numpy as np
import tttrlib
from dataclasses import dataclass, field
from typing import List, Tuple, Dict, Optional, Union, Any


@dataclass
class Burst:
    start: int
    end: int  # inclusive
    peak_score: float


@dataclass
class KalmanBurstResult:
    x_filt: np.ndarray  # (T, dim) filtered rates
    P_filt: np.ndarray  # (T, dim, dim) filtered covariances
    innovation_mahal: np.ndarray  # (T,) Mahalanobis innovation distances D_t
    z_thresh: float
    bursts: List[Burst]


@dataclass
class KalmanBurstDetector:
    dim: int
    dt: float
    q: float = 1.0
    r_scale: float = 1.0
    z_thresh: float = 5.0
    min_len: int = 5
    merge_gap: int = 1
    x0: Optional[np.ndarray] = None
    P0: Optional[np.ndarray] = None

    _A: np.ndarray = field(init=False, repr=False)
    _Q: np.ndarray = field(init=False, repr=False)
    _H: np.ndarray = field(init=False, repr=False)

    def __post_init__(self):
        self._A = np.eye(self.dim)
        self._H = np.eye(self.dim)
        self._Q = np.eye(self.dim) * self.q
        if self.x0 is None:
            self.x0 = np.zeros(self.dim)
        if self.P0 is None:
            self.P0 = np.eye(self.dim) * 1e6

    def _adaptive_R(self, y_rate: np.ndarray, x_pred: np.ndarray) -> np.ndarray:
        rate = np.maximum(x_pred, 1e-12)
        var_rate = rate / self.dt
        R = np.diag(self.r_scale * var_rate)
        return R

    def detect(self, counts: np.ndarray) -> KalmanBurstResult:
        counts = np.asarray(counts, dtype=float)
        assert counts.ndim == 2 and counts.shape[1] == self.dim, "counts must be of shape (T, dim)"
        T = counts.shape[0]
        y = counts / self.dt

        x = self.x0.copy()
        P = self.P0.copy()

        x_filt = np.zeros((T, self.dim))
        P_filt = np.zeros((T, self.dim, self.dim))
        D_mahal = np.zeros(T)

        I = np.eye(self.dim)

        for t in range(T):
            x_pred = self._A @ x
            P_pred = self._A @ P @ self._A.T + self._Q
            R = self._adaptive_R(y[t], x_pred)
            v = y[t] - self._H @ x_pred
            S = self._H @ P_pred @ self._H.T + R
            try:
                S_inv = np.linalg.inv(S)
            except np.linalg.LinAlgError:
                S_inv = np.linalg.pinv(S)
            K = P_pred @ self._H.T @ S_inv
            x = x_pred + K @ v
            P = (I - K @ self._H) @ P_pred
            x_filt[t] = x
            P_filt[t] = P
            D_mahal[t] = float(np.sqrt(v.T @ S_inv @ v))

        bursts = self._extract_bursts(D_mahal)

        return KalmanBurstResult(x_filt=x_filt, P_filt=P_filt, innovation_mahal=D_mahal, z_thresh=self.z_thresh, bursts=bursts)

    def _extract_bursts(self, D: np.ndarray) -> List[Burst]:
        over = D > self.z_thresh
        bursts: List[Burst] = []
        t = 0
        T = over.size
        while t < T:
            if over[t]:
                start = t
                peak = D[t]
                while t + 1 < T and over[t + 1]:
                    t += 1
                    if D[t] > peak:
                        peak = D[t]
                end = t
                if end - start + 1 >= self.min_len:
                    bursts.append(Burst(start, end, peak))
            t += 1

        if self.merge_gap > 0 and len(bursts) > 1:
            merged: List[Burst] = [bursts[0]]
            for b in bursts[1:]:
                last = merged[-1]
                if b.start - last.end - 1 <= self.merge_gap:
                    merged[-1] = Burst(start=last.start, end=b.end, peak_score=max(last.peak_score, b.peak_score))
                else:
                    merged.append(b)
            bursts = merged

        return bursts


def bin_photons(donor: np.ndarray, acceptor: np.ndarray, dt: float = 1e-3, tmax: Optional[float] = None) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Bin photon timestamps into count arrays.
    
    Args:
        donor: Array of donor photon timestamps in seconds
        acceptor: Array of acceptor photon timestamps in seconds
        dt: Bin width in seconds
        tmax: Maximum time to consider. If None, uses the maximum timestamp.
        
    Returns:
        Tuple of (D_counts, A_counts, bin_edges) where:
        - D_counts is an array of donor counts per bin
        - A_counts is an array of acceptor counts per bin
        - bin_edges is an array of bin edges
    """
    if tmax is None:
        tmax = max(donor.max() if len(donor) > 0 else 0, 
                  acceptor.max() if len(acceptor) > 0 else 0)
    
    n_bins = int(np.ceil(tmax / dt))
    
    D_counts, bin_edges = np.histogram(donor, bins=n_bins, range=(0, tmax))
    A_counts, _ = np.histogram(acceptor, bins=n_bins, range=(0, tmax))
    
    return D_counts, A_counts, bin_edges


def bin_photons_multi(timestamps_list: List[np.ndarray], dt: float = 1e-3, tmax: Optional[float] = None) -> Tuple[np.ndarray, np.ndarray]:
    """
    Bin photon timestamps from multiple channels into count arrays.
    
    Args:
        timestamps_list: List of arrays containing photon timestamps in seconds for each channel
        dt: Bin width in seconds
        tmax: Maximum time to consider. If None, uses the maximum timestamp across all channels.
        
    Returns:
        Tuple of (counts, bin_edges) where:
        - counts is a 2D array of shape (T, n_channels) containing counts per bin for each channel
        - bin_edges is an array of bin edges
    """
    if not timestamps_list:
        return np.array([]), np.array([])
    
    # Find maximum timestamp across all channels
    if tmax is None:
        tmax = 0
        for timestamps in timestamps_list:
            if len(timestamps) > 0:
                tmax = max(tmax, timestamps.max())
    
    n_bins = int(np.ceil(tmax / dt))
    n_channels = len(timestamps_list)
    
    # Initialize counts array
    counts = np.zeros((n_bins, n_channels), dtype=np.int64)
    
    # Bin timestamps for each channel
    for i, timestamps in enumerate(timestamps_list):
        if len(timestamps) > 0:
            counts[:, i], bin_edges = np.histogram(timestamps, bins=n_bins, range=(0, tmax))
        else:
            bin_edges = np.linspace(0, tmax, n_bins + 1)
    
    return counts, bin_edges


def kalman_burst_detection(
    donor_timestamps: np.ndarray, 
    acceptor_timestamps: np.ndarray,
    dt: float = 1e-3,
    q: float = 20.0,
    r_scale: float = 1.0,
    z_thresh: float = 3.0,
    min_len: int = 2,
    merge_gap: int = 20,
    min_counts: int = 20
) -> Tuple[List[Dict[str, Any]], np.ndarray, np.ndarray, np.ndarray, KalmanBurstResult]:
    """
    Perform Kalman filter-based burst detection on donor and acceptor photon timestamps.
    
    Args:
        donor_timestamps: Array of donor photon timestamps in seconds
        acceptor_timestamps: Array of acceptor photon timestamps in seconds
        dt: Bin width in seconds
        q: Process noise parameter
        r_scale: Measurement noise scaling parameter
        z_thresh: Threshold for burst detection
        min_len: Minimum burst length in bins
        merge_gap: Maximum gap between bursts to merge them
        min_counts: Minimum number of photons (donor + acceptor) for a burst
        
    Returns:
        Tuple of (bursts, bin_edges, filtered_rates, innovation_distances, result) where:
        - bursts is a list of dictionaries with keys 'start_bin', 'end_bin', 'peak_score'
        - bin_edges is an array of bin edges
        - filtered_rates is a (T, 2) array of filtered rates [donor_rate, acceptor_rate]
        - innovation_distances is a (T,) array of Mahalanobis innovation distances
        - result is the KalmanBurstResult object
    """
    # Bin the photon timestamps
    D_counts, A_counts, bin_edges = bin_photons(donor_timestamps, acceptor_timestamps, dt)
    
    # Stack counts for Kalman filter
    counts = np.stack([D_counts, A_counts], axis=1)
    
    # Create and run the Kalman filter detector
    detector = KalmanBurstDetector(
        dim=2,
        dt=dt,
        q=q,
        r_scale=r_scale,
        z_thresh=z_thresh,
        min_len=min_len,
        merge_gap=merge_gap
    )
    
    result = detector.detect(counts)
    
    # Convert bursts to the format expected by the rest of the code
    bursts = []
    for b in result.bursts:
        # Calculate the total counts in this burst
        start_bin = b.start
        end_bin = b.end
        donor_counts = np.sum(counts[start_bin:end_bin+1, 0])
        acceptor_counts = np.sum(counts[start_bin:end_bin+1, 1])
        total_counts = donor_counts + acceptor_counts
        
        # Skip bursts with too few photons
        if total_counts < min_counts:
            continue
            
        bursts.append({
            'start_bin': start_bin,
            'end_bin': end_bin,
            'peak_score': b.peak_score,
            'start_time': bin_edges[start_bin],
            'end_time': bin_edges[end_bin + 1] if end_bin + 1 < len(bin_edges) else bin_edges[-1],
            'duration': bin_edges[end_bin + 1] - bin_edges[start_bin] if end_bin + 1 < len(bin_edges) else bin_edges[-1] - bin_edges[start_bin],
            'donor': donor_counts,
            'acceptor': acceptor_counts,
            'FRET': acceptor_counts / total_counts if total_counts > 0 else np.nan
        })
    
    return bursts, bin_edges, result.x_filt, result.innovation_mahal, result


def kalman_burst_detection_multi(
    timestamps_list: List[np.ndarray],
    dt: float = 1e-3,
    q: float = 20.0,
    r_scale: float = 1.0,
    z_thresh: float = 3.0,
    min_len: int = 2,
    merge_gap: int = 20,
    min_counts: int = 20
) -> Tuple[List[Dict[str, Any]], np.ndarray, np.ndarray, np.ndarray, KalmanBurstResult]:
    """
    Perform Kalman filter-based burst detection on multiple channels.
    
    Args:
        timestamps_list: List of arrays containing photon timestamps in seconds for each channel
        dt: Bin width in seconds
        q: Process noise parameter
        r_scale: Measurement noise scaling parameter
        z_thresh: Threshold for burst detection
        min_len: Minimum burst length in bins
        merge_gap: Maximum gap between bursts to merge them
        min_counts: Minimum number of photons (across all channels) for a burst
        
    Returns:
        Tuple of (bursts, bin_edges, filtered_rates, innovation_distances, result) where:
        - bursts is a list of dictionaries with keys 'start_bin', 'end_bin', 'peak_score', 'counts'
        - bin_edges is an array of bin edges
        - filtered_rates is a (T, n_channels) array of filtered rates
        - innovation_distances is a (T,) array of Mahalanobis innovation distances
        - result is the KalmanBurstResult object
    """
    # Check if we have any channels
    if not timestamps_list:
        return [], np.array([]), np.array([]), np.array([]), None
    
    # Bin the photon timestamps
    counts, bin_edges = bin_photons_multi(timestamps_list, dt)
    
    # Get the number of channels
    n_channels = counts.shape[1]
    
    # Create and run the Kalman filter detector
    detector = KalmanBurstDetector(
        dim=n_channels,
        dt=dt,
        q=q,
        r_scale=r_scale,
        z_thresh=z_thresh,
        min_len=min_len,
        merge_gap=merge_gap
    )
    
    result = detector.detect(counts)
    
    # Convert bursts to the format expected by the rest of the code
    bursts = []
    for b in result.bursts:
        # Calculate the total counts in this burst
        start_bin = b.start
        end_bin = b.end
        
        # Get counts for each channel in this burst
        channel_counts = np.sum(counts[start_bin:end_bin+1, :], axis=0)
        total_counts = np.sum(channel_counts)
        
        # Skip bursts with too few photons
        if total_counts < min_counts:
            continue
        
        # Create a burst dictionary with counts for each channel
        burst_dict = {
            'start_bin': start_bin,
            'end_bin': end_bin,
            'peak_score': b.peak_score,
            'start_time': bin_edges[start_bin],
            'end_time': bin_edges[end_bin + 1] if end_bin + 1 < len(bin_edges) else bin_edges[-1],
            'duration': bin_edges[end_bin + 1] - bin_edges[start_bin] if end_bin + 1 < len(bin_edges) else bin_edges[-1] - bin_edges[start_bin],
            'counts': channel_counts.tolist(),  # Store counts for each channel
            'total_counts': total_counts
        }
        
        # For backward compatibility, if we have exactly 2 channels, add donor/acceptor fields
        if n_channels == 2:
            burst_dict['donor'] = channel_counts[0]
            burst_dict['acceptor'] = channel_counts[1]
            burst_dict['FRET'] = channel_counts[1] / total_counts if total_counts > 0 else np.nan
            
        bursts.append(burst_dict)
    
    return bursts, bin_edges, result.x_filt, result.innovation_mahal, result


def create_burst_mask(bursts: List[Dict[str, Any]], tttr_length: int) -> np.ndarray:
    """
    Create a boolean mask for the TTTR data based on the bursts.
    
    Args:
        bursts: List of burst dictionaries with 'start_time' and 'end_time' keys
        tttr_length: Length of the TTTR data
        
    Returns:
        Boolean mask array of length tttr_length
    """
    mask = np.zeros(tttr_length, dtype=bool)
    
    for burst in bursts:
        start_time = burst['start_time']
        end_time = burst['end_time']
        
        # Find indices in the TTTR data that fall within this burst
        # This is a placeholder - actual implementation would depend on how
        # the TTTR data is structured and how to map from time to indices
        
    return mask


def convert_bursts_to_start_stop(bursts: List[Dict[str, Any]], tttr: tttrlib.TTTR) -> np.ndarray:
    """
    Convert bursts to start-stop indices for the TTTR data.
    
    Args:
        bursts: List of burst dictionaries with 'start_time' and 'end_time' keys
        tttr: TTTR object
        
    Returns:
        Array of shape (n_bursts, 2) with [start_idx, end_idx] for each burst
    """
    if not bursts:
        return np.array([], dtype=np.uint64).reshape(0, 2)
    
    macro_times = tttr.macro_times
    time_unit = tttr.header.macro_time_resolution
    
    start_stop = []
    for burst in bursts:
        start_time = burst['start_time']
        end_time = burst['end_time']
        
        # Convert times to macro time units
        start_mt = int(start_time / time_unit)
        end_mt = int(end_time / time_unit)
        
        # Find indices in the macro_times array
        start_idx = np.searchsorted(macro_times, start_mt)
        end_idx = np.searchsorted(macro_times, end_mt, side='right') - 1
        
        if start_idx <= end_idx:
            start_stop.append([start_idx, end_idx])
    
    if not start_stop:
        return np.array([], dtype=np.uint64).reshape(0, 2)
    
    return np.array(start_stop, dtype=np.uint64)