import numpy as np
from scipy.optimize import linprog
import quantecon.game_theory as gt

"""
Explanation: 零和ゲームのナッシュ均衡
quantecon.game_theory と scipy.optimize.linprog で求めてみる
End of explanation
"""


U = np.array(
    [[0, -1, 1],
     [1, 0, -1],
     [-1, 1, 0]]
)

"""
Explanation: じゃんけん・ゲームを例に計算してみる．
あとあと便利なので，NumPy array としてプレイヤー0の利得行列を定義しておく:
End of explanation
"""


p0 = gt.Player(U)
p1 = gt.Player(-U.T)

"""
Explanation: quantecon.game_theory でナッシュ均衡を求める
End of explanation
"""


g = gt.NormalFormGame((p0, p1))

print(g)

gt.lemke_howson(g)

gt.support_enumeration(g)

"""
Explanation: プレイヤー1の行列は -U の転置 (.T) であることに注意．
End of explanation
"""


m, n = U.shape

"""
Explanation: scipy.optimize.linprog で線形計画問題を解くことでナッシュ均衡を求める
主問題:
$$
\min u
$$
subject to
$$
U y - \mathbf{1} u \leq \mathbf{0},\quad
\mathbf{1}' y = 1,\quad
y \geq \mathbf{0}
$$
これを
scipy.optimize.linprog
の形式
$$
\max c' x
$$
subject to
$$
A_{\mathit{ub}} x \leq b_{\mathit{ub}},\quad
A_{\mathit{eq}} x = b_{\mathit{eq}},\quad
l \leq x \leq u
$$
に合わせると，
$$
x =
\begin{pmatrix}
y \ u
\end{pmatrix},\ %
c =
\begin{pmatrix}
\mathbf{0} \ 1
\end{pmatrix},\ %
A_{\mathit{ub}} =
\begin{pmatrix}
U & -\mathbf{1}
\end{pmatrix},\ %
b_{\mathit{ub}} = \mathbf{0},\ %
A_{\mathit{eq}} =
\begin{pmatrix}
\mathbf{1}' & 0
\end{pmatrix},\ %
b_{\mathit{eq}} = \begin{pmatrix}1\end{pmatrix}
$$
$$
l =
\begin{pmatrix}
0 & \cdots & 0 & -\infty
\end{pmatrix}',\ %
u =
\begin{pmatrix}
\infty & \cdots & \infty
\end{pmatrix}'
$$
両プレイヤーの戦略の数をそれぞれ $m$, $n$ とする:
End of explanation
"""


c = np.zeros(n+1)
c[-1] = 1
c

A_ub = np.hstack((U, np.full((m, 1), -1)))
A_ub

b_ub = np.zeros(m)
b_ub

A_eq = np.ones((1, n+1))
A_eq[0, -1] = 0
A_eq

b_eq = np.ones(1)
b_eq

bounds = [(0, None)] * n + [(None, None)]
bounds

"""
Explanation: 各入力を定義する:
End of explanation
"""


res_p = linprog(c, A_ub, b_ub, A_eq, b_eq, bounds)

"""
Explanation: scipy.optimize.linprog に渡して解かせる:
End of explanation
"""


res_p

"""
Explanation: 結果:
End of explanation
"""


res_p.x[:-1]

"""
Explanation: プレイヤー1の均衡戦略:
End of explanation
"""


res_p.x[-1]

"""
Explanation: ゲームの値:
End of explanation
"""


c = np.zeros(m+1)
c[-1] = -1

A_ub = np.hstack((-U.T, np.full((n, 1), 1)))
b_ub = np.zeros(n)

A_eq = np.ones((1, m+1))
A_eq[0, -1] = 0
b_eq = np.ones(1)

bounds = [(0, None)] * m + [(None, None)]

res_d = linprog(c, A_ub, b_ub, A_eq, b_eq, bounds)
res_d

"""
Explanation: scipy.optimize.linprog は双対解を返してくれないようなので，双対問題も定式化して解かせる:
$$
\min -v
$$
subject to
$$
-U' x + \mathbf{1} v \leq \mathbf{0},\quad
\mathbf{1}' x = 1,\quad
x \geq \mathbf{0}
$$
End of explanation
"""


res_d.x[:-1]

"""
Explanation: プレイヤー0の均衡戦略:
End of explanation
"""


res_d.x[-1]

"""
Explanation: ゲームの値:
End of explanation
"""


res_p.x[-1] - res_d.x[-1]

"""
Explanation: 微妙は誤差はある:
End of explanation
"""


def solve_zerosum_lemke_howson(U):
    g = gt.NormalFormGame((gt.Player(U), gt.Player(-U.T)))
    NE = gt.lemke_howson(g)
    return NE

def solve_zerosum_linprog(U, method='revised simplex'):
    U = np.asarray(U)
    m, n = U.shape
    
    # Primal problem
    c = np.zeros(n+1)
    c[-1] = 1
    A_ub = np.hstack((U, np.full((m, 1), -1)))
    b_ub = np.zeros(m)
    A_eq = np.ones((1, n+1))
    A_eq[0, -1] = 0
    b_eq = np.ones(1)
    bounds = [(0, None)] * n + [(None, None)]
    
    res_p = linprog(c, A_ub, b_ub, A_eq, b_eq, bounds, method=method)
    
    # Dual problem
    c = np.zeros(m+1)
    c[-1] = -1
    A_ub = np.hstack((-U.T, np.full((n, 1), 1)))
    b_ub = np.zeros(n)
    A_eq = np.ones((1, m+1))
    A_eq[0, -1] = 0
    b_eq = np.ones(1)
    bounds = [(0, None)] * m + [(None, None)]
    
    res_d = linprog(c, A_ub, b_ub, A_eq, b_eq, bounds, method=method)
    
    NE = (res_d.x[:-1], res_p.x[:-1])
    
    return NE

solve_zerosum_lemke_howson(U)

solve_zerosum_linprog(U)

"""
Explanation: 関数としてまとめてみる
End of explanation
"""


m, n = 4, 3
U = np.random.randn(m, n)
U

solve_zerosum_lemke_howson(U)

solve_zerosum_linprog(U)

"""
Explanation: ランダムに行列を発生させて計算させてみる
End of explanation
"""


