import numpy as np
from scipy.stats import chi2
from collections import Counter
import pandas as pd
from itertools import product

def chi2_randomness(sequence, s, alpha=0.05):
    """
    Chi-square test for the presence of Markovian memory in a discrete sequence
    Args:
        sequence: Discrete sequence 
        s: Order of the Markov chain (memory depth)
        alpha: Significance level
        
    Returns:
        Dict: Statistic, critical value, p-value, etc.
    """
    n = len(sequence) # the number of elements in the sequence under study
    if n <= s:
        return {
            "statistic": 0.0, "critical_value": 0.0, "p_value": 1.0,
            "is_rejected": False, "df": 0, "warning": "Sequence too short"
        }
    
    alphabet = sorted(list(set(sequence))) # automatically determines the system's alphabet (set E)
    N = len(alphabet) # number of elements in E
    
    # 1. We count only the frequencies of the (s+1)-blocks (v_ij), N^{s+1} combinations
    v_ij_counts = Counter() # create a Counter object
    for t in range(n - s):
    # We traverse the entire history using a window of length s+1. Each combination found is stored as a tuple.
        chain_s_plus_1 = tuple(sequence[t : t + s + 1]) # create a key from (s+1)-blocks
        v_ij_counts[chain_s_plus_1] += 1 # If the pattern is encountered for the first time, the Counter object will create a new record
        
    # 2. To obtain the frequency v_i, we sum the frequencies v_ij over all j for a given prefix i
    v_i_from_v_ij = Counter() # create a Counter object
    for chain_s_plus_1, count in v_ij_counts.items():
        prefix_i = chain_s_plus_1[:-1] # create a new key; take the first s elements
        v_i_from_v_ij[prefix_i] += count # calculate the frequency v_i

    # 3. Calculate the chi-square statistic 
    chi_square_stat = 0
    
    # We iterate through all encountered s-length state sequences (i) in a loop
    for i_chain, v_i in v_i_from_v_ij.items():       
        expected = v_i / N # Expected count under H0: v_i / N (equal probability)
        # For each i-sequence, we consider all possible continuations j from the alphabet (set E)
        for j_state in alphabet:
            ij_chain = i_chain + (j_state,)
            v_ij = v_ij_counts.get(ij_chain, 0)             
            chi_square_stat += ((v_ij - expected)**2) / expected
   
    df = (N**s) * (N - 1) # Degrees of freedom: N^s(N - 1)
    critical_value = chi2.ppf(1 - alpha, df) 
    p_value = 1 - chi2.cdf(chi_square_stat, df)
    
    return {
        "statistic": chi_square_stat,
        "critical_value": critical_value,
        "p_value": p_value,
        "is_rejected": chi_square_stat > critical_value, 
        "df": df
    }
    
def quantile_transform(data, n_states=4)-> np.ndarray:
    """
    Quantile transformation divides real-valued data 
    into N discrete states based on empirical quantiles.
    This neutralizes trend effects and ensures P{E(j)} = 1/N
    
    Args:
        data: One-dimensional array of real numbers (increments, returns)
        n_states: Number of discrete states (alphabet size)
        
    Returns:
        np.ndarray: Sequence of integers (state indices) from 0 to N-1
    
    """
    # Calculate the thresholds for N states
    # For n_states=4, the thresholds will be 25%, 50%, and 75%
    quantiles = np.linspace(0, 100, n_states + 1)[1:-1]
    thresholds = np.percentile(data, quantiles)
    
    # We convert the data into a sequence of states 0, 1, ..., N-1
    return np.digitize(data, thresholds) 
    
def analyze_patterns(sequence, s_order):
    """
    Analyzes the deviation of the frequencies of blocks of length s+1 from the expected count
    """
    n = len(sequence)
    s_plus_1 = s_order + 1
    
    if n < s_plus_1:
        print(f"Warning: Sequence length ({n}) is too short for order s={s_order}")
        return pd.DataFrame() 
    
    alphabet = sorted(list(set(sequence)))
    N = len(alphabet)
    
    # 1. Collect all blocks of length s+1
    chains = [tuple(sequence[t : t + s_plus_1]) for t in range(n - s_order)] # list of tuples
    counts = Counter(chains) # Counter converts a list into a dictionary where the key is a pattern and the value is the number of times it appeared
    total_observed = len(chains)
    
    # 2. Expected count for ideal randomness 
    num_combinations = N**s_plus_1
    expected_count = total_observed / num_combinations # Evi_j = n/N^(s+1)
    
    results = []
    # 3. Generate the Cartesian product of the alphabet (all possible state sequences)
    for pattern in product(alphabet, repeat=s_plus_1): # For example, E × E × E = E³
        actual = counts.get(pattern, 0)
        # Calculate the percentage deviation from the expected count under H0
        diff_pct = ((actual - expected_count) / expected_count) * 100 if expected_count > 0 else 0
        
        results.append({
            "Pattern": pattern,
            "Actual": actual,
            "Expected": round(expected_count, 1),
            "Deviation_%": round(diff_pct, 2)
        })
    
    return pd.DataFrame(results)     
    
if __name__ == "__main__":
   
    # --- 1: price increments with a positive expected value (trend) ---
    print("\n1. price increments with a positive expected value (trend):")   
    data = np.random.normal(loc=0.5, scale=1.0, size=1000)
    
    # TEST 1: Split by zero threshold 
    seq_raw = (data > 0).astype(int)
    res_raw = chi2_randomness(seq_raw, s=1)
    
    # TEST 2: Median Split 
    seq_quant = quantile_transform(data, n_states=2)
    res_quant = chi2_randomness(seq_quant, s=1)   
    counts_q = Counter(seq_quant)    
    print(f"  Without quantiles: X2={res_raw['statistic']:.2f}, P-value={res_raw['p_value']:.4f}")
    print(f"  With quantiles (balance {counts_q[0]}/{counts_q[1]}): X2={res_quant['statistic']:.2f}, P-value={res_quant['p_value']:.4f}")

    # --- 2: Pendulum Markov chain model) ---
    print("\n2. pendulum Markov chain model:")
    markov_seq = [0]
    for _ in range(999):
        curr = markov_seq[-1]
        # The probability of a state change is 90%
        markov_seq.append(np.random.choice([0, 1], p=[0.1, 0.9]) if curr == 0 else np.random.choice([0, 1], p=[0.9, 0.1]))
    
    res_markov = chi2_randomness(markov_seq, s=1)
    counts = Counter(markov_seq)
    
    print(f"  Distribution: 0: {counts[0]}, 1: {counts[1]}")
    print(f"  Test result: X2={res_markov['statistic']:.2f}, P-value={res_markov['p_value']:.4f}")