import networkx as nx
import matplotlib.pyplot as plt

# 1. 定义状态和转移规则
# (起点, 终点, 概率)
transitions = [
    (1, 1, 0.5), # 状态1 50% 回到自己
    (1, 2, 0.5), # 状态1 50% 去状态2
    (2, 3, 1.0), # 状态2 100% 去状态3
    (3, 3, 1.0)  # 状态3 100% 留在大结局
]

# 2. 创建图对象
G = nx.DiGraph()
for u, v, p in transitions:
    G.add_edge(u, v, weight=p)

# 3. 设置绘图布局 (让它们排成一行)
pos = {1: (0, 0), 2: (1, 0), 3: (2, 0)}

plt.figure(figsize=(10, 4))

# 4. 绘制节点
nx.draw_networkx_nodes(G, pos, node_size=2000, node_color='lightblue', edgecolors='black')
nx.draw_networkx_labels(G, pos, font_size=20)

# 5. 绘制边 (箭头)
# 直线边
nx.draw_networkx_edges(G, pos, edgelist=[(1, 2), (2, 3)], arrowstyle='->', arrowsize=20)
# 自环边 (弧形)
nx.draw_networkx_edges(G, pos, edgelist=[(1, 1), (3, 3)], connectionstyle='arc3, rad=0.5', arrowstyle='->', arrowsize=20)

# 6. 添加概率标签
plt.text(0, 0.25, "0.5", ha='center', fontsize=12, color='red') # 1->1
plt.text(0.5, 0.05, "0.5", ha='center', fontsize=12, color='red') # 1->2
plt.text(1.5, 0.05, "1.0", ha='center', fontsize=12, color='red') # 2->3
plt.text(2, 0.25, "1.0", ha='center', fontsize=12, color='red') # 3->3

plt.axis('off')
plt.title("Markov Chain Visualization", fontsize=16)
plt.show()

import networkx as nx
import matplotlib.pyplot as plt

def draw_chain(transitions, title, ax):
    G = nx.DiGraph()
    # 添加边和权重
    for u, v, p in transitions:
        G.add_edge(u, v, weight=p)

    # 使用圆形布局
    pos = nx.circular_layout(G)
    
    # 画节点
    nx.draw_networkx_nodes(G, pos, ax=ax, node_size=2000, node_color='lightblue', edgecolors='black')
    nx.draw_networkx_labels(G, pos, ax=ax, font_size=16, font_weight='bold')

    # 画边 (带弧度，避免重叠)
    nx.draw_networkx_edges(G, pos, ax=ax, arrowstyle='->', arrowsize=25, 
                           connectionstyle='arc3, rad=0.15')
    
    # 标注概率
    edge_labels = {(u, v): f"{p}" for u, v, p in transitions}
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, ax=ax, label_pos=0.3, font_size=12)

    ax.set_title(title, fontsize=14)
    ax.axis('off')

# --- 1. 定义可约链 (Reducible) ---
# 这里的 "3" 是个死胡同 (吸收态)，回不到 "1" 或 "2"
transitions_reducible = [
    (1, 1, 0.5), 
    (1, 2, 0.5), 
    (2, 3, 1.0), 
    (3, 3, 1.0) 
]

# --- 2. 定义不可约链 (Irreducible) ---
# 我们添加了 3->1，形成闭环，所有状态互通
transitions_irreducible = [
    (1, 1, 0.5), 
    (1, 2, 0.5), 
    (2, 3, 1.0), 
    (3, 1, 1.0) 
]

# --- 绘图 ---
fig, axes = plt.subplots(1, 2, figsize=(12, 5))

draw_chain(transitions_reducible, "Reducible (Broken Flow)", axes[0])
draw_chain(transitions_irreducible, "Irreducible (Connected Flow)", axes[1])

plt.tight_layout()
plt.show()

import networkx as nx
import matplotlib.pyplot as plt

def draw_subplot(transitions, title, ax, pos_type='circular'):
    G = nx.DiGraph()
    for u, v in transitions:
        G.add_edge(u, v)
    
    if pos_type == 'circular':
        pos = nx.circular_layout(G)
    else:
        pos = nx.spring_layout(G, seed=42) # 固定种子以保持形状
        
    # 绘制节点和边
    nx.draw_networkx_nodes(G, pos, ax=ax, node_size=800, node_color='lightgray', edgecolors='black')
    nx.draw_networkx_labels(G, pos, ax=ax, font_size=12, font_weight='bold')
    
    # 绘制带箭头的边 (使用arc3弯曲，防止重叠)
    nx.draw_networkx_edges(G, pos, ax=ax, arrowstyle='->', arrowsize=20, 
                           connectionstyle='arc3, rad=0.1')
    
    # 检查是否有自环 (Self-loop)，手动标记一下方便观察
    self_loops = list(nx.selfloop_edges(G))
    if self_loops:
        nx.draw_networkx_edges(G, pos, ax=ax, edgelist=self_loops, 
                               arrowstyle='->', arrowsize=20, connectionstyle='arc3, rad=0.5')

    ax.set_title(title, fontsize=10)
    ax.axis('off')

# --- 定义六种情况 ---

# 1. 简单的周期性 (Periodic)
# 最简单的一来一回，周期为 2
t_periodic = [(1, 2), (2, 1)]

# 2. 简单的非周期性 (Aperiodic)
# 只要有自环，立刻打破周期，周期为 1
t_aperiodic = [(1, 1)]

# 3. 不可约 & 周期性 (Irreducible & Periodic)
# 一个大环，大家互通（不可约），但步数必须是 3 的倍数（周期 3）
t_irr_per = [(1, 2), (2, 3), (3, 1)]

# 4. 不可约 & 非周期性 (Irreducible & Aperiodic)
# 大家互通（不可约），但在 1 处加了个自环，节奏乱了（非周期）
t_irr_aper = [(1, 2), (2, 3), (3, 1), (1, 1)]

# 5. 可约 & 周期性 (Reducible & Periodic)
# 1 去 2 回不来（可约）。2 和 3 互跳（周期 2）。
# 注意：这里指 recurrent class {2,3} 是周期的。
t_red_per = [(1, 2), (2, 3), (3, 2)]

# 6. 可约 & 非周期性 (Reducible & Aperiodic)
# 1 去 2 回不来（可约）。2 有自环（非周期）。
t_red_aper = [(1, 2), (2, 2)]

# --- 绘图 ---
fig, axes = plt.subplots(2, 3, figsize=(15, 8))
axes = axes.flatten()

draw_subplot(t_periodic, "1. Periodic (General)\n(Period=2)", axes[0])
draw_subplot(t_aperiodic, "2. Aperiodic (General)\n(Self-loop)", axes[1])
draw_subplot(t_irr_per, "3. Irreducible & Periodic\n(Cycle=3)", axes[2])
draw_subplot(t_irr_aper, "4. Irreducible & Aperiodic\n(Cycle=3 + Self-loop)", axes[3])
draw_subplot(t_red_per, "5. Reducible & Periodic\n(1->{2,3}, {2,3} is periodic)", axes[4])
draw_subplot(t_red_aper, "6. Reducible & Aperiodic\n(1->2, 2 is aperiodic)", axes[5])

plt.tight_layout()
plt.show()

import numpy as np
import networkx as nx
import matplotlib.pyplot as plt

# 1. 定义一个无向图 (Undirected Graph)
# 这里的权重 (weight) 可以理解为两个状态之间的“亲密度”或“通道宽窄”
G = nx.Graph()
G.add_edge(0, 1, weight=2)
G.add_edge(1, 2, weight=1)

# 2. 自动构建转移矩阵 P
# 规则：从节点 i 跳到邻居 j 的概率 = (i,j的权重) / (i 的所有权重之和)
nodes = sorted(G.nodes())
n = len(nodes)
P = np.zeros((n, n))

for i in nodes:
    total_weight = sum([G[i][nbr]['weight'] for nbr in G.neighbors(i)])
    for j in G.neighbors(i):
        P[i, j] = G[i][j]['weight'] / total_weight

print("--- 转移矩阵 P ---")
print(P)

# 3. 计算稳态分布 pi
# 技巧：对于无向图，pi_i 正比于节点 i 的“度” (所有连接边的权重和)
degrees = [sum([G[i][nbr]['weight'] for nbr in G.neighbors(i)]) for i in nodes]
total_degree_sum = sum(degrees)
pi = np.array([d / total_degree_sum for d in degrees])

print("\n--- 稳态分布 pi ---")
print(f"pi = {pi}")

# 4. 验证细致平衡 (Detailed Balance): pi_i * P_ij = pi_j * P_ji ?
print("\n--- 验证细致平衡 (Flow Check) ---")
# 检查 状态 0 <-> 状态 1
flow_0_to_1 = pi[0] * P[0, 1]
flow_1_to_0 = pi[1] * P[1, 0]

print(f"Flow 0 -> 1: {pi[0]:.4f} * {P[0, 1]:.4f} = {flow_0_to_1:.4f}")
print(f"Flow 1 -> 0: {pi[1]:.4f} * {P[1, 0]:.4f} = {flow_1_to_0:.4f}")

if np.isclose(flow_0_to_1, flow_1_to_0):
    print("✅ 状态 0 和 1 之间满足细致平衡！")
else:
    print("❌ 不平衡")

# --- 绘图 ---
pos = nx.spring_layout(G, seed=42)
plt.figure(figsize=(6, 4))
nx.draw(G, pos, with_labels=True, node_color='lightblue', node_size=2000, font_size=16)
labels = nx.get_edge_attributes(G, 'weight')
nx.draw_networkx_edge_labels(G, pos, edge_labels=labels)
plt.title("Reversible Markov Chain (Undirected Graph)")
plt.show()

import numpy as np
import matplotlib.pyplot as plt

# ============ 基础工具 ============

def is_row_stochastic(P, tol=1e-12):
    """检查转移矩阵是否“行随机”（每行和约为1，元素>=0）"""
    P = np.asarray(P, dtype=float)
    nonneg = np.all(P >= -tol) # 检查非负性
    rowsum_one = np.allclose(P.sum(axis=1), 1.0, atol=1e-10) # 检查每行和是否约为1
    return bool(nonneg and rowsum_one), P.sum(axis=1)

def n_step_transition(P, n):
    """n 步转移矩阵：P^n"""
    return np.linalg.matrix_power(np.asarray(P, dtype=float), n)

def simulate_markov_chain(P, init_state, n_steps, rng=None):
    """
    从单个初始状态模拟一条马尔可夫链路径。
    P: 行随机矩阵；init_state: int (0..S-1)；返回数组 shape=(n_steps+1,)
    """
    if rng is None:
        rng = np.random.default_rng()
    P = np.asarray(P, dtype=float)
    S = P.shape[0] # 状态数
    path = np.empty(n_steps+1, dtype=int) # 初始化路径
    path[0] = int(init_state) # 确保初始状态是整数
    for t in range(n_steps): # 逐步生成路径
        i = path[t]
        path[t+1] = rng.choice(S, p=P[i]) # 从当前状态 i 选择下一个状态
    return path

def simulate_many(P, pi0, n_steps, n_runs=10000, rng=None):
    """
    模拟多条路径，估计各时刻的经验分布（与理论 pi0 P^t 对比）。
    返回：
      emp_dist: shape=(n_steps+1, S) 经验分布
      th_dist : shape=(n_steps+1, S) 理论分布
    """
    if rng is None:
        rng = np.random.default_rng()
    P = np.asarray(P, dtype=float)
    S = P.shape[0]
    # 理论分布随时间演化
    th = np.zeros((n_steps+1, S))
    th[0] = pi0
    for t in range(n_steps):
        th[t+1] = th[t] @ P
    # 经验分布
    counts = np.zeros((n_steps+1, S), dtype=int)
    init_states = rng.choice(S, size=n_runs, p=pi0)
    for r in range(n_runs):
        s0 = init_states[r]
        path = simulate_markov_chain(P, s0, n_steps, rng=rng)
        for t in range(n_steps+1):
            counts[t, path[t]] += 1
    emp = counts / n_runs
    return emp, th

# ============ 示例 1：两状态（晴/雨） ============

# 状态编码：0=晴(S), 1=雨(R)
P2 = np.array([[0.8, 0.2],
               [0.4, 0.6]], dtype=float)

ok, rowsums = is_row_stochastic(P2)
print("P2 行随机性检查：", ok, " 行和=", rowsums)

# 单一路径模拟与可视化
rng = np.random.default_rng(2025)
path = simulate_markov_chain(P2, init_state=0, n_steps=50, rng=rng)

plt.figure(figsize=(9,3))
plt.plot(range(len(path)), path, marker='o')
plt.hlines(0, -1, len(path)-1, colors='green', linestyles='dashed', label="Sunny")
plt.hlines(1, -1, len(path)-1, colors='yellow', linestyles='dashed', label="Rainy")
plt.text(10, -0.4, "样本路径（状态随时间跳变的折线）", fontsize=12, color='red')
plt.title("Sample path of 2-state Markov chain (0=Sunny, 1=Rainy)")
plt.xlabel("time t")
plt.ylabel("state")
plt.legend()
plt.show()

# 多路径统计 vs 理论分布
pi0 = np.array([0.5, 0.5])  # 初始分布
emp, th = simulate_many(P2, pi0, n_steps=20, n_runs=5000, rng=rng)

# 画“处于 Sunny 的概率”随时间变化：理论 vs 经验
plt.figure(figsize=(9,4))
plt.plot(th[:,0], label="theory P(X_t=Sunny)")
plt.plot(emp[:,0], label="empirical P(X_t=Sunny)")
plt.text(3, 0.6, "理论分布与模拟经验分布几乎重合（随样本数增大吻合更好）。", fontsize=12, color='red')
plt.title("Distribution evolution in 2-state chain")
plt.xlabel("time t")
plt.ylabel("probability of Sunny")
plt.legend()
plt.show()

# n 步转移矩阵示例
P2_5 = n_step_transition(P2, 5)
print("P2^5 =\n", np.round(P2_5, 4))

# ============ 示例 2：三状态（含吸收态 C） ============

# 状态编码：0=A, 1=B, 2=C(吸收)
P3 = np.array([[0.6, 0.4, 0.0],
               [0.2, 0.5, 0.3],
               [0.0, 0.0, 1.0]], dtype=float)

ok3, rowsums3 = is_row_stochastic(P3)
print("P3 行随机性检查：", ok3, " 行和=", rowsums3)

path3 = simulate_markov_chain(P3, init_state=0, n_steps=40, rng=rng)

plt.figure(figsize=(9,3))
plt.plot(range(len(path3)), path3, marker='o')
plt.text(10, 1.5, "样本路径（状态随时间跳变的折线）\n状态 C 一旦进入就不离开（吸收态）。", fontsize=12, color='red')
plt.title("Sample path of 3-state chain (2 is absorbing)")
plt.xlabel("time t")
plt.ylabel("state (0=A, 1=B, 2=C)")
plt.show()

# 多路径统计：观测吸收到 C 的概率随时间变化
pi0_3 = np.array([1.0, 0.0, 0.0])  # 从 A 起步
emp3, th3 = simulate_many(P3, pi0_3, n_steps=20, n_runs=5000, rng=rng)

plt.figure(figsize=(9,4))
plt.plot(th3[:,2], label="theory P(X_t=C)")
plt.plot(emp3[:,2], label="empirical P(X_t=C)")
plt.text(5, 0.2, "理论分布与模拟经验分布几乎重合（随样本数增大吻合更好）。\n吸收态 C 的概率随时间单调上升", fontsize=12, color='red')
plt.title("Absorption probability into state C over time")
plt.xlabel("time t")
plt.ylabel("probability in C")
plt.legend()
plt.show()

# n 步转移矩阵示例
P3_5 = n_step_transition(P3, 5)
print("P3^5 =\n", np.round(P3_5, 4))

import numpy as np
import matplotlib.pyplot as plt

P = np.array([[0.9, 0.1],
              [0.5, 0.5]])

# 初始分布
mu = np.array([1.0, 0.0])  

distributions = [mu]
for _ in range(20):
    mu = mu @ P
    distributions.append(mu)

distributions = np.array(distributions)

plt.plot(distributions[:,0], label="Pr[state=0]")
plt.plot(distributions[:,1], label="Pr[state=1]")
plt.axhline(0.833, color="gray", linestyle="--", label="π0")
plt.axhline(0.167, color="gray", linestyle="--", label="π1")
plt.text(2.5, 0.7, "从状态 0 出发，概率分布逐步收敛到 (0.833, 0.167)", fontsize=12, color='red')
plt.xlabel("Step")
plt.ylabel("Probability")
plt.legend()
plt.title("Convergence to Stationary Distribution")
plt.show()

import numpy as np
import matplotlib.pyplot as plt

# 三状态马尔可夫链转移矩阵
P = np.array([[0.5, 0.5, 0.0],
              [0.2, 0.5, 0.3],
              [0.0, 0.3, 0.7]])

# 初始分布（全部在状态0）
mu = np.array([1.0, 0.0, 0.0])

# 计算平稳分布：解 pi P = pi
eigvals, eigvecs = np.linalg.eig(P.T)
stat_dist = eigvecs[:, np.isclose(eigvals, 1)]
stat_dist = stat_dist[:,0]
stat_dist = stat_dist / stat_dist.sum()  # 归一化
print("Stationary distribution:", stat_dist.real)

# 迭代分布演化
distributions = [mu]
for _ in range(30):
    mu = mu @ P
    distributions.append(mu)

distributions = np.array(distributions)

# 绘图
plt.figure(figsize=(8,5))
for i in range(3):
    plt.plot(distributions[:, i], label=f"Pr[state={i}]")
    plt.axhline(stat_dist[i].real, linestyle="--", color="gray")

plt.text(5, 0.8, "从初始状态全在 0 出发，\n随着步数增加，概率分布逐渐收敛到这组平稳分布；\n横虚线表示平稳分布值。", fontsize=12, color='red')
plt.xlabel("Step")
plt.ylabel("Probability")
plt.title("Convergence of 3-State Markov Chain")
plt.legend()
plt.show()

import numpy as np
import matplotlib.pyplot as plt

def ring_rw_transition_matrix(n):
    """生成 n 状态的环形随机游走转移矩阵"""
    P = np.zeros((n, n))
    for i in range(n):
        P[i, (i-1)%n] = 0.5
        P[i, (i+1)%n] = 0.5
    return P

def simulate_chain(P, steps=30, start_state=0):
    """从单点分布开始，计算分布演化"""
    n = P.shape[0]
    mu = np.zeros(n)
    mu[start_state] = 1.0
    distributions = [mu]
    for _ in range(steps):
        mu = mu @ P
        distributions.append(mu)
    return np.array(distributions)

# 参数
steps = 30
P5 = ring_rw_transition_matrix(5)
P6 = ring_rw_transition_matrix(6)

# 模拟
dist5 = simulate_chain(P5, steps)
dist6 = simulate_chain(P6, steps)

# 平稳分布（对于奇数 n，均匀分布；偶数情况不存在唯一收敛）
pi5 = np.ones(5) / 5
pi6 = np.ones(6) / 6

# 绘图
fig, axes = plt.subplots(2, 1, figsize=(8,8), sharex=True)

# n=5
for i in range(5):
    axes[0].plot(dist5[:, i], label=f"state {i}")
axes[0].hlines(pi5, 0, steps, colors="gray", linestyles="--", linewidth=1)
axes[0].set_title("Ring Random Walk (n=5, odd → converges to uniform)")
axes[0].set_ylabel("Probability")
axes[0].legend()

# n=6
for i in range(6):
    axes[1].plot(dist6[:, i], label=f"state {i}")
axes[1].hlines(pi6, 0, steps, colors="gray", linestyles="--", linewidth=1)
axes[1].set_title("Ring Random Walk (n=6, even → oscillates)")
axes[1].set_xlabel("Step")
axes[1].set_ylabel("Probability")
axes[1].legend()

plt.tight_layout()
plt.show()

# Simulate Markov chains and compute total variation distance (TV) to the stationary distribution.
# We will:
# 1. Define several transition matrices (fast/slow 3-state, cycle random walks n=5 and n=6).
# 2. Compute TV distance over time starting from state 0.
# 3. Compute mixing times tau(epsilon) for epsilons = [0.1, 0.01, 0.001].
# 4. Plot TV vs time for comparisons and show a table of mixing times.
#
# Note: Plots use matplotlib (no seaborn) and each figure is a single plot as requested.

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

def stationary_from_P(P):
    # solve pi = pi P  with sum(pi)=1  -> transpose eigenvector of P^T for eigenvalue 1
    w, v = np.linalg.eig(P.T)
    idx = np.argmin(np.abs(w - 1.0))
    pi = np.real(v[:, idx])
    pi = pi / pi.sum()
    pi = np.maximum(pi, 0)
    pi = pi / pi.sum()
    return pi

def tv_distance(p, q):
    return 0.5 * np.sum(np.abs(p - q))

def tv_curve(P, p0, t_max):
    n = P.shape[0]
    pis = stationary_from_P(P)
    p = p0.copy()
    tvs = []
    for t in range(t_max + 1):
        tvs.append(tv_distance(p, pis))
        p = p @ P
    return np.array(tvs), pis

def mixing_time_from_tvs(tvs, eps):
    # minimal t such that tvs[t] <= eps
    below = np.where(tvs <= eps)[0]
    return int(below[0]) if below.size > 0 else np.nan

# Define chains
# 3-state fast chain
P_fast = np.array([
    [0.6, 0.3, 0.1],
    [0.2, 0.6, 0.2],
    [0.1, 0.3, 0.6]
])

# 3-state slow chain (more "sticky" on state 0)
P_slow = np.array([
    [0.9, 0.08, 0.02],
    [0.2, 0.7, 0.1],
    [0.15, 0.15, 0.7]
])

# Cycle random walk
def cycle_P(n):
    P = np.zeros((n, n))
    for i in range(n):
        P[i, (i+1) % n] = 0.5
        P[i, (i-1) % n] = 0.5
    return P

P_cycle5 = cycle_P(5)
P_cycle6 = cycle_P(6)

# initial distribution: start at state 0
def e0(n): 
    v = np.zeros(n); v[0]=1.0; return v

t_max = 200

# compute tv curves
tvs_fast, pi_fast = tv_curve(P_fast, e0(3), t_max)
tvs_slow, pi_slow = tv_curve(P_slow, e0(3), t_max)

tvs_c5, pi_c5 = tv_curve(P_cycle5, e0(5), t_max)
tvs_c6, pi_c6 = tv_curve(P_cycle6, e0(6), t_max)

# compute mixing times for selected epsilons
epsilons = [1e-1, 1e-2, 1e-3]
rows = []
for name, tvs in [
    ("3-state fast", tvs_fast),
    ("3-state slow", tvs_slow),
    ("cycle n=5", tvs_c5),
    ("cycle n=6", tvs_c6)
]:
    entry = {"chain": name}
    for eps in epsilons:
        entry[f"tau({eps})"] = mixing_time_from_tvs(tvs, eps)
    rows.append(entry)

df_mix = pd.DataFrame(rows)

# Plot 1: 3-state fast vs slow
plt.figure(figsize=(8,4))
plt.plot(tvs_fast, label='3-state fast')
plt.plot(tvs_slow, label='3-state slow')
plt.yscale('log')  # show both fast and slow clearly on log scale
plt.xlabel('time t')
plt.ylabel('TV distance (log scale)')
plt.title('Total Variation distance: 3-state fast vs slow (start at state 0)')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()

# Plot 2: cycle n=5 vs n=6
plt.figure(figsize=(8,4))
plt.plot(tvs_c5, label='cycle n=5')
plt.plot(tvs_c6, label='cycle n=6')
plt.yscale('log')
plt.xlabel('time t')
plt.ylabel('TV distance (log scale)')
plt.title('Total Variation distance: cycle random walk n=5 vs n=6 (start at state 0)')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()


# Also print stationary distributions for reference
pi_table = pd.DataFrame({
    "chain": ["3-state fast", "3-state slow", "cycle n=5", "cycle n=6"],
    "stationary": [pi_fast, pi_slow, pi_c5, pi_c6]
})
pi_table['stationary_str'] = pi_table['stationary'].apply(lambda x: np.array2string(x, precision=4, separator=', '))
pi_table = pi_table[['chain','stationary_str']]
display("Stationary distributions", pi_table)