import numpy as np
import matplotlib.pyplot as plt

# Parameters
beta = 0.95          # Discount factor
grid_size = 100
grid_max = 10
x_grid = np.linspace(0, grid_max, grid_size)

# Reward function: utility from consumption
def u(c):
    return np.log(c + 1e-10)  # Add epsilon to avoid log(0)

# Bellman operator
def bellman_operator(V, beta=0.95):
    V_new = np.zeros_like(V)
    for i, x in enumerate(x_grid):
        # Consumption choice: consume fraction of current state
        c_vals = np.linspace(0, x, grid_size)
        x_prime = x - c_vals  # Next period state
        # Interpolate V at x_prime
        V_interp = np.interp(x_prime, x_grid, V)
        # Maximize over consumption
        objective = u(c_vals) + beta * V_interp
        V_new[i] = np.max(objective)
    return V_new

# Iterate to convergence
V = np.zeros(grid_size)
for iteration in range(100):
    V_new = bellman_operator(V, beta)
    if np.max(np.abs(V_new - V)) < 1e-6:
        print(f"Converged in {iteration} iterations")
        break
    V = V_new

print(f"Final value function: V(0)={V[0]:.4f}, V(max)={V[-1]:.4f}")

# Extract policy function
policy = np.zeros(grid_size)
for i, x in enumerate(x_grid):
    c_vals = np.linspace(0, x, grid_size)
    x_prime = x - c_vals
    V_interp = np.interp(x_prime, x_grid, V)
    objective = u(c_vals) + beta * V_interp
    policy[i] = c_vals[np.argmax(objective)]

print(f"Policy at x=1: consume {policy[10]:.3f}")
print(f"Policy at x=5: consume {policy[50]:.3f}")
print(f"Policy at x=10: consume {policy[-1]:.3f}")

import numpy as np

# Define a simple 3-state Markov chain
# States: {low, medium, high}
P = np.array([
    [0.7, 0.2, 0.1],   # From low
    [0.2, 0.6, 0.2],   # From medium
    [0.1, 0.3, 0.6]    # From high
])

# Find stationary distribution via eigenvalue decomposition
eigenvalues, eigenvectors = np.linalg.eig(P.T)
stationary_idx = np.argmax(np.abs(eigenvalues - 1.0) < 1e-10)
stationary = eigenvectors[:, stationary_idx].real
stationary = stationary / stationary.sum()

print("Transition matrix P:")
print(P)
print(f"\nStationary distribution: {stationary}")
print(f"Sum: {stationary.sum():.6f}")

# Simulate Markov chain
np.random.seed(42)
n_periods = 1000
state = 0  # Start in low state
path = [state]

for _ in range(n_periods - 1):
    state = np.random.choice(3, p=P[state, :])
    path.append(state)

# Empirical distribution
unique, counts = np.unique(path, return_counts=True)
empirical = counts / n_periods

print(f"Empirical distribution after {n_periods} periods:")
for state, freq in zip(unique, empirical):
    print(f"  State {state}: {freq:.3f}")
print(f"\nCompare to stationary: {stationary}")

import numpy as np
from scipy.optimize import minimize

# Generate synthetic AR(1) data
np.random.seed(42)
true_rho = 0.8
true_sigma = 1.0
n_obs = 200

y = np.zeros(n_obs)
for t in range(1, n_obs):
    y[t] = true_rho * y[t-1] + np.random.normal(0, true_sigma)

# Log-likelihood function
def neg_log_likelihood(params):
    rho, sigma = params
    if sigma <= 0 or abs(rho) >= 1:
        return 1e10  # Penalty for invalid parameters

    residuals = y[1:] - rho * y[:-1]
    ll = -0.5 * n_obs * np.log(2 * np.pi * sigma**2)
    ll -= 0.5 * np.sum(residuals**2) / sigma**2
    return -ll

# Optimize
result = minimize(neg_log_likelihood, x0=[0.5, 1.0], method='Nelder-Mead')
rho_hat, sigma_hat = result.x

print(f"True parameters: rho={true_rho:.3f}, sigma={true_sigma:.3f}")
print(f"MLE estimates:  rho={rho_hat:.3f}, sigma={sigma_hat:.3f}")
print(f"Estimation error: rho={abs(rho_hat-true_rho):.4f}, sigma={abs(sigma_hat-true_sigma):.4f}")

# QuantEcon example (pseudo-code, library may not be installed)
# from quantecon import MarkovChain
#
# mc = MarkovChain(P)
# print(mc.stationary_distributions)
# print(mc.simulate(ts_length=1000, init=0))

print("QuantEcon provides:")
print("- MarkovChain: stationary distributions, simulation")
print("- DynamicProgram: value iteration, policy iteration")
print("- ARMA: time series modeling")
print("- Estimation: GMM, MLE, Kalman filtering")
print("\nInstall: pip install quantecon")

using QuantEcon

# Markov chain in Julia
P = [0.7 0.2 0.1;
     0.2 0.6 0.2;
     0.1 0.3 0.6]

mc = MarkovChain(P)
stationary_distributions(mc)
simulate(mc, 1000, init=1)