"""
Markov Chain Models for Economic Simulation

This module implements the core Markov chain models representing:
1. Capitalist Chain (M-C-M'): Money → Commodities → More Money
2. Consumer Chain (C-M-C): Commodities → Money → Commodities

Based on Marx's economic theory and Piketty's inequality principle (r > g).
"""

import numpy as np
from typing import List, Dict, Tuple
import random


class MarkovChain:
    """
    Base class for Markov chain implementations.
    
    Provides the foundation for economic state transitions with configurable
    transition probabilities and state tracking.
    """
    
    def __init__(self, states: List[str], transition_probability: float, initial_state: str):
        """
        Initialize the Markov chain.
        
        Args:
            states: List of possible states
            transition_probability: Base probability for state transitions
            initial_state: Starting state for the chain
        """
        self.states = states
        self.transition_probability = max(0.0, min(1.0, transition_probability))
        self.current_state = initial_state
        self.state_history = [initial_state]
        self.iteration_count = 0
        
        # Validate initial state
        if initial_state not in states:
            raise ValueError(f"Initial state '{initial_state}' not in states list")
    
    def get_transition_matrix(self) -> np.ndarray:
        """
        Get the transition probability matrix for this chain.
        Must be implemented by subclasses.
        
        Returns:
            numpy array representing the transition matrix
        """
        raise NotImplementedError("Subclasses must implement get_transition_matrix")
    
    def step(self) -> str:
        """
        Perform one step in the Markov chain.
        
        Returns:
            The new current state after the transition
        """
        current_index = self.states.index(self.current_state)
        transition_matrix = self.get_transition_matrix()
        probabilities = transition_matrix[current_index]
        
        # Choose next state based on probabilities
        next_state_index = np.random.choice(len(self.states), p=probabilities)
        self.current_state = self.states[next_state_index]
        
        self.state_history.append(self.current_state)
        self.iteration_count += 1
        
        return self.current_state
    
    def simulate(self, iterations: int) -> List[str]:
        """
        Run the Markov chain for multiple iterations.
        
        Args:
            iterations: Number of steps to simulate
            
        Returns:
            List of states visited during simulation
        """
        for _ in range(iterations):
            self.step()
        
        return self.state_history.copy()
    
    def get_state_distribution(self) -> Dict[str, float]:
        """
        Calculate the current state distribution from history.
        
        Returns:
            Dictionary mapping states to their frequency ratios
        """
        if not self.state_history:
            return {state: 0.0 for state in self.states}
        
        total_count = len(self.state_history)
        distribution = {}
        
        for state in self.states:
            count = self.state_history.count(state)
            distribution[state] = count / total_count
            
        return distribution
    
    def reset(self, initial_state: str = None):
        """
        Reset the chain to initial conditions.
        
        Args:
            initial_state: New initial state (optional)
        """
        if initial_state and initial_state in self.states:
            self.current_state = initial_state
        else:
            self.current_state = self.state_history[0]
            
        self.state_history = [self.current_state]
        self.iteration_count = 0


class CapitalistChain(MarkovChain):
    """
    Represents the capitalist economic cycle: M-C-M' (Money → Commodities → More Money)
    
    This chain models capital accumulation where money is invested in commodities
    to generate more money, with returns based on the capital rate (r).
    """
    
    def __init__(self, capital_rate: float):
        """
        Initialize the capitalist chain.
        
        Args:
            capital_rate: The rate of return on capital (r)
        """
        states = ['Money', 'Commodities', 'Enhanced_Money']
        super().__init__(states, capital_rate, 'Money')
        self.capital_rate = capital_rate
        self.wealth_multiplier = 1.0
    
    def get_transition_matrix(self) -> np.ndarray:
        """
        Create transition matrix for M-C-M' cycle.
        
        The matrix ensures:
        - Money transitions to Commodities with probability r
        - Commodities always transition to Enhanced_Money (probability 1)
        - Enhanced_Money always transitions back to Money (probability 1)
        
        Returns:
            3x3 transition probability matrix
        """
        r = self.transition_probability
        
        # States: ['Money', 'Commodities', 'Enhanced_Money']
        matrix = np.array([
            [1-r,   r,    0.0],  # Money -> stay or go to Commodities
            [0.0,   0.0,  1.0],  # Commodities -> Enhanced_Money (always)
            [1.0,   0.0,  0.0]   # Enhanced_Money -> Money (always)
        ])
        
        return matrix
    
    def step(self) -> str:
        """
        Perform one step and update wealth when completing a cycle.
        
        Returns:
            The new current state
        """
        previous_state = self.current_state
        new_state = super().step()
        
        # If we completed a full cycle (Enhanced_Money -> Money), increase wealth
        if previous_state == 'Enhanced_Money' and new_state == 'Money':
            self.wealth_multiplier *= (1 + self.capital_rate)
        
        return new_state
    
    def get_current_wealth(self) -> float:
        """
        Get the current accumulated wealth.
        
        Returns:
            Current wealth multiplier representing accumulated capital
        """
        return self.wealth_multiplier


class ConsumerChain(MarkovChain):
    """
    Represents the consumer economic cycle: C-M-C (Commodities → Money → Commodities)
    
    This chain models consumption patterns where commodities are exchanged for money
    to purchase other commodities, growing with the economic growth rate (g).
    """
    
    def __init__(self, growth_rate: float):
        """
        Initialize the consumer chain.
        
        Args:
            growth_rate: The economic growth rate (g)
        """
        states = ['Commodities', 'Money', 'New_Commodities']
        super().__init__(states, growth_rate, 'Commodities')
        self.growth_rate = growth_rate
        self.consumption_value = 1.0
    
    def get_transition_matrix(self) -> np.ndarray:
        """
        Create transition matrix for C-M-C cycle.
        
        The matrix ensures:
        - Commodities transition to Money with probability g
        - Money always transitions to New_Commodities (probability 1)
        - New_Commodities always transition back to Commodities (probability 1)
        
        Returns:
            3x3 transition probability matrix
        """
        g = self.transition_probability
        
        # States: ['Commodities', 'Money', 'New_Commodities']
        matrix = np.array([
            [1-g,   g,    0.0],  # Commodities -> stay or go to Money
            [0.0,   0.0,  1.0],  # Money -> New_Commodities (always)
            [1.0,   0.0,  0.0]   # New_Commodities -> Commodities (always)
        ])
        
        return matrix
    
    def step(self) -> str:
        """
        Perform one step and update consumption value when completing a cycle.
        
        Returns:
            The new current state
        """
        previous_state = self.current_state
        new_state = super().step()
        
        # If we completed a full cycle (New_Commodities -> Commodities), grow consumption
        if previous_state == 'New_Commodities' and new_state == 'Commodities':
            self.consumption_value *= (1 + self.growth_rate)
        
        return new_state
    
    def get_current_consumption(self) -> float:
        """
        Get the current consumption value.
        
        Returns:
            Current consumption value representing accumulated consumption capacity
        """
        return self.consumption_value


class EconomicAgent:
    """
    Represents an individual economic agent with both capitalist and consumer behaviors.
    
    Each agent operates both chains simultaneously, allowing for mixed economic activity
    and wealth accumulation patterns.
    """
    
    def __init__(self, agent_id: str, capital_rate: float, growth_rate: float, 
                 initial_capital: float = 1000.0, initial_consumption: float = 1000.0):
        """
        Initialize an economic agent.
        
        Args:
            agent_id: Unique identifier for the agent
            capital_rate: Capital return rate (r)
            growth_rate: Economic growth rate (g)
            initial_capital: Starting capital amount
            initial_consumption: Starting consumption capacity
        """
        self.agent_id = agent_id
        self.capitalist_chain = CapitalistChain(capital_rate)
        self.consumer_chain = ConsumerChain(growth_rate)
        
        self.initial_capital = initial_capital
        self.initial_consumption = initial_consumption
        
        # Track wealth over time
        self.wealth_history = []
        self.consumption_history = []
    
    def step(self) -> Tuple[float, float]:
        """
        Perform one simulation step for both chains.
        
        Returns:
            Tuple of (current_wealth, current_consumption)
        """
        # Step both chains
        self.capitalist_chain.step()
        self.consumer_chain.step()
        
        # Calculate current values
        current_wealth = self.initial_capital * self.capitalist_chain.get_current_wealth()
        current_consumption = self.initial_consumption * self.consumer_chain.get_current_consumption()
        
        # Store history
        self.wealth_history.append(current_wealth)
        self.consumption_history.append(current_consumption)
        
        return current_wealth, current_consumption
    
    def get_total_wealth(self) -> float:
        """
        Get the agent's total economic value.
        
        Returns:
            Sum of capital wealth and consumption capacity
        """
        if not self.wealth_history or not self.consumption_history:
            return self.initial_capital + self.initial_consumption
            
        return self.wealth_history[-1] + self.consumption_history[-1]
    
    def get_wealth_ratio(self) -> float:
        """
        Get the ratio of capital wealth to total wealth.
        
        Returns:
            Ratio representing the capitalist vs consumer wealth proportion
        """
        total = self.get_total_wealth()
        if total == 0:
            return 0.0
            
        if not self.wealth_history:
            return self.initial_capital / (self.initial_capital + self.initial_consumption)
            
        return self.wealth_history[-1] / total