#-*- encoding: utf-8 -*-
%matplotlib inline
# 日本語対応
import matplotlib
matplotlib.rcParams['font.family'] = 'Osaka'
import math
import numpy as np
import pandas as pd
import functools
import matplotlib.pyplot as plt
import scipy.stats as stats
np.set_printoptions(precision=3)
np.set_printoptions(linewidth=400)
np.set_printoptions(threshold=np.nan)
pd.set_option('display.max_columns', 30)
pd.set_option('display.width', 400)
pd.set_option('display.precision', 3)

"""
Explanation: Appendix1  ゲームの期数に関して
実験本体の詳細:  README
End of explanation
"""


# settings
seed = 282
rs = np.random.RandomState(seed)
parameter = 0.03

# plot
fig = plt.figure(figsize=(20, 12))
mu = 33.333
sigma = 1077.778

# theoretical pmf
ax = plt.subplot(2, 1, 1)
plt.title("Geo(0.03)の確率関数")
x = np.arange(1, 301)
plt.plot(x, stats.geom.pmf(x, parameter), linewidth=1, color='green', label="theoretical pmf")
plt.xlabel("ts_length(X)")
plt.ylabel("probability(f)")
ax.text(35, 0.02, r'''$\mu$={0:.3f}, $\sigma^2$={1:.3f}'''.format(mu, sigma), ha = 'left', va = 'bottom', size=14)
ax.axvline(x=mu, linewidth=1, color='red', label="average")
ax.set_xlim([1, 300])
plt.legend()

# theoretical cmf
ax = plt.subplot(2, 1, 2)
plt.title("Geo(0.03)の下側分布関数")
x = np.arange(1, 301)
plt.plot(x, stats.geom.cdf(x, parameter), linewidth=1, color='blue', label="theoretical cdf")
plt.xlabel("ts_length(X)")
plt.ylabel("cumulative dens(F)")
ax.text(35, 0.55, r'''$\mu$={0:.3f}のときF={1:.3f}'''.format(mu, stats.geom.cdf(mu, parameter)), ha = 'left', va = 'bottom', size=14)
ax.axvline(x=mu, linewidth=1, color='red', label="average")
ax.set_xlim([1, 300])
plt.legend()

plt.show()

"""
Explanation: 各セッションの期数の決め方は、以下の通り。
引用: README#共通ルール
<blockquote>
<p>ゲームを何期続けるかは、以下のルールによって決めます。

<ol>
<li>第1期は確率1で到来する
<li>以降は毎期末に抽選を行い、**97%の確率でゲームを継続する**（3%の確率で終了する）  
（これは、無限回繰り返しゲームにおいて現在割引価値を0.97と取ることを意味します）<br>
</ol>

<p>平均は33.33期になります。
</blockquote>

したがって、各セッションの期数を確率変数 $X_1, X_2, \ldots, X_{1000}$ とすると、
$$
X_1, X_2, \ldots, X_{1000} \overset{\text{i.i.d.}}{\sim} Geo(0.03)
$$
となる。ただしGeo(p)は「成功確率がpの時、初めて成功するまでの総試行回数」を表す幾何分布で、（下側）累積分布関数は
$$
P(X \leq x) = F(x) = \sum_{k=1}^{x} p(1-p)^{k-1} \hspace{20pt} x=1, 2, \ldots
$$
となる。
平均は
\begin{eqnarray}
E[X] &=& \sum_{x=1}^{\infty} x pq^{x-1}  \hspace{20pt} q =  1-p\
&=& p\left(\frac{\partial}{\partial q} \sum_{x=1}^{\infty} q^x \right) \
&=& p \left(\frac{\partial}{\partial q} \frac{q}{1-q} \right) \
&=& p \frac{1-q+q}{(1-q)^2} \
&=& \frac{1}{p}
\end{eqnarray}
分散は
\begin{eqnarray}
E[X(X+1)] &=& \sum_{x=1}^{\infty} x(x+1) pq^{x-1}\
&=& p\left(\frac{\partial^2}{\partial q^2} \sum_{x=1}^{\infty} q^{x+1} \right) \
&=& p \left(\frac{\partial^2}{\partial q^2} \frac{q^2}{1-q} \right) \
&=& p \left(\frac{\partial}{\partial q} \frac{2q(1-q)+q^2}{(1-q)^2} \right) \
&=& p \left(\frac{\partial}{\partial q} \frac{2q-q^2}{(1-q)^2} \right) \
&=& p \left(\frac{(2-2q)(1-q)^2 + 2(2q-q^2)(1-q)}{(1-q)^4} \right) \
&=& p \left(\frac{2(1-q)^2 + 4q-2q^2)}{(1-q)^3} \right) \
&=& \frac{2}{p^2}
\end{eqnarray}
\begin{eqnarray}
Var(X) &=& E[X^2] - (E[X])^2 \
&=& E[X(X+1)] - E[X]  - (E[X])^2 \
&=& \frac{2}{p^2} - \frac{1}{p} - \frac{1}{p^2} \
&=& \frac{1-p}{p^2}
\end{eqnarray}
となる。
今回はp=0.03なので、$E[X] = 33.33, Var(X) = 1077.78$ となる。
確率関数と累積分布関数
確率（質量）関数と下側累積分布関数をplot（理論値）。線にしてあるが、実際は離散値（x=1, 2, 3, ...）。
End of explanation
"""


seed = 282
np.random.seed(seed=seed)
rs = np.random.RandomState(seed)
parameter = 0.03
size = 1000
sum_max = 100000
ave = 1/parameter
var = (1-parameter)/parameter**2
sum_x = np.arange(100001)

# 「幾何分布から1000個のサンプルを取り出し、和を取る」
#  ことを10万回繰り返し、シミュレーションで分布を求める
sim_size = 100000
sums = np.zeros(sim_size, dtype=float)
for i in range(sim_size):
    sums[i] = rs.geometric(p=parameter, size=1000).sum()

# 負の二項分布 n=1000, p=0.03
# scipy.statの負の二項分布は「失敗回数」を数える分布なので、適宜修正
nbin = stats.nbinom.pmf(sum_x-1000, 1000, parameter)

# 正規分布 mu=3333.33, sigma^2=107777.78
normal_ave = size * ave
normal_var = size * var
normal = stats.norm.pdf(x=sum_x, loc=normal_ave, scale=pow(normal_var, 0.5))

# プロット
fig, ax = plt.subplots(figsize=(20, 6))
plt.title("Geo(0.03)から抽出した1000個のサンプルの和S(X)の分布")
plt.plot(sum_x, nbin, linewidth=2, color='blue', label="theoretical pmf(負の二項分布)")
plt.plot(sum_x, normal, linewidth=2, color='orange', label="approx pdf(正規分布)")
plt.hist(sums, bins=100, color='#cccccc', normed=True, label="simulation(size=100000)")
ax.set_xlim([25000, 45000])
plt.xlabel("ts_length(X)")
plt.ylabel("density(f)")
plt.legend()
plt.show()

"""
Explanation: 確率関数は単調減少。$P(X \leq E[X]) = F(33.33) = 0.634$ より、全体の63.4%程度が平均よりも小さい値であることが期待される。
このことは、幾何分布を適当な指数分布で近似（連続化）できるとすれば、その時間あたりの生起回数はポアソン分布$F(x) = \sum_{k=0}^x\frac{\lambda^k}{k!}e^{- \lambda}$に従うことからもわかる。今1単位時間を33.33期とすれば、1単位時間に1回以上「ゲームが終了する」事象が起こる確率は、$\lambda = 1$ のポアソン分布で$x \geq 1$となる確率なので、
\begin{eqnarray}
P(X\geq1) = 1 - F(0) = \frac{1^0}{0!}e^{- 1} = \frac{1}{e} = 0.6321
\end{eqnarray}
となって、ほぼ一致する。
和の分布
正確な和の分布
幾何分布に従う確率変数 $X_1, X_2, \ldots, X_{1000}$ の和を $S(X)$ とすると、
$$
S(X) \ {\sim} \ Negbin(n=1000, p=0.03)
$$
となる。ただしNegbin(n, p)は「成功確率がpの時、n回成功するまでの総試行回数」を表す負の二項分布で、累積分布関数は、
$$
P(Y \leq y) = F(y) = \sum_{k=n}^{y} {}{k-1}C {n-1} p^n (1-p)^{k-n} \hspace{20pt} y=n, n+1, n+2, \ldots
$$
となる。平均は $E[Y] = \frac{n}{p}$、分散は $Var(Y) = \frac{n(1-p)}{p^2}$ となる。
いま $p=0.03, n=1000$ なので、$E[S(X)] = 33333.33, Var(S(X)) = 1077777.78$ となる。
和の近似的な分布
$Var(X_i) < \infty$ より、 $X_1, X_2, \ldots, X_n$ の平均 $Y = E[X]$ の分布は、
\begin{eqnarray}
P(Y\leq y) = F(y) \overset{\text{d}}{\longrightarrow} \int_{-\infty}^{\infty} \frac{\sqrt{n}}{\sqrt{2 \pi \sigma^2}} exp\left( {-\cfrac{n(y-\mu)^2}{2\sigma^2}} \right ) \ dy
\end{eqnarray}
となって、正規分布 $N(\mu, \frac{\sigma^2}{n})$ に分布収束する。
したがって和の近似的な分布は、（$Z = nY$ と変数変換すれば良いので）
$$
S(X) = n E[X] \sim N(n\mu, n\sigma^2)
$$
となる。
$X_i$ の母平均, 母分散は $\mu = 0.03, \sigma = 1077.77\ldots$ だったので、$E[S] = 33.33, Var(S) = 1077777.78$ となる。
和の分布の①正確分布、②近似分布、③シミュレーションを比較
End of explanation
"""


# 統計量
data = np.zeros((11, 3), dtype=float)

ordered_sum = np.sort(sums)
sim_lower_percent = lambda p: (ordered_sum[sim_size*p-1]+ordered_sum[sim_size*p])/2
data[0, 0] = sums.mean()
data[1, 0] = sums.var()
data[2, 0] = sim_lower_percent(0.01)
data[3, 0] = sim_lower_percent(0.025)
data[4, 0] = sim_lower_percent(0.05)
data[5, 0] = sim_lower_percent(0.1)
data[6, 0] = np.median(sums)
data[7, 0] = sim_lower_percent(0.9)
data[8, 0] = sim_lower_percent(0.95)
data[9, 0] = sim_lower_percent(0.975)
data[10, 0] = sim_lower_percent(0.99)

data[0, 1] = stats.nbinom.mean(1000, parameter) + 1000
data[1, 1] = stats.nbinom.var(1000, parameter)
data[2, 1] = stats.nbinom.interval(0.98, 1000, parameter)[0]+1000
data[3, 1] = stats.nbinom.interval(0.95, 1000, parameter)[0]+1000
data[4, 1] = stats.nbinom.interval(0.9, 1000, parameter)[0]+1000
data[5, 1] = stats.nbinom.interval(0.8, 1000, parameter)[0]+1000
data[6, 1] = stats.nbinom.median(1000, parameter) + 1000
data[7, 1] = stats.nbinom.interval(0.8, 1000, parameter)[1]+1000
data[8, 1] = stats.nbinom.interval(0.9, 1000, parameter)[1]+1000
data[9, 1] = stats.nbinom.interval(0.95, 1000, parameter)[1]+1000
data[10, 1] = stats.nbinom.interval(0.98, 1000, parameter)[1]+1000

data[0, 2] = stats.norm.mean(loc=normal_ave, scale=pow(normal_var, 0.5))
data[1, 2] = stats.norm.var(loc=normal_ave, scale=pow(normal_var, 0.5))
data[2, 2] = stats.norm.interval(0.98, loc=normal_ave, scale=pow(normal_var, 0.5))[0]
data[3, 2] = stats.norm.interval(0.95, loc=normal_ave, scale=pow(normal_var, 0.5))[0]
data[4, 2] = stats.norm.interval(0.9, loc=normal_ave, scale=pow(normal_var, 0.5))[0]
data[5, 2] = stats.norm.interval(0.8, loc=normal_ave, scale=pow(normal_var, 0.5))[0]
data[6, 2] = stats.norm.median(loc=normal_ave, scale=pow(normal_var, 0.5))
data[7, 2] = stats.norm.interval(0.8, loc=normal_ave, scale=pow(normal_var, 0.5))[1]
data[8, 2] = stats.norm.interval(0.9, loc=normal_ave, scale=pow(normal_var, 0.5))[1]
data[9, 2] = stats.norm.interval(0.95, loc=normal_ave, scale=pow(normal_var, 0.5))[1]
data[10, 2] = stats.norm.interval(0.98, loc=normal_ave, scale=pow(normal_var, 0.5))[1]

df = pd.DataFrame(data,
    columns=["シミュレーション", "正確分布", "近似分布"],
    index=["平均", "分散", "下側 1%", "下側 2.5%", "下側 5%", "下側 10%", "中央値", "上側 10%", "上側 5%", "上側 2.5%", "上側 1%"])
print(df)

"""
Explanation: 近似分布の方が正確な分布よりもやや右側に寄っているが、ほぼ一致している。細かく統計量を見ると、
End of explanation
"""


# 総期数の仮説検定
def two_sided_p(x):
    parameter = 0.03
    med = stats.nbinom.median(1000, parameter) + 1000
    if x <= med:
        for p in np.arange(0, 1.005, 0.005):
            point = stats.nbinom.interval(p, 1000, parameter)[0]+1000
            if point <= x:
                return 1-p 
    else:
        for p in np.arange(0, 1.005, 0.005):
            point = stats.nbinom.interval(p, 1000, parameter)[1]+1000
            if x <= point:
                return 1-p
            
    raise ValueError


def one_sided_p(x):
    parameter = 0.03
    med = stats.nbinom.median(1000, parameter)+1000
    if x <= med:
        p = stats.nbinom.cdf(x-1000, 1000, parameter)
        return p
    else:
        p = stats.nbinom.cdf(x-1000, 1000, parameter)
        return 1-p 

    raise ValueError


print("尾山ゼミの総期数 S=32856")
print("両側検定でのp値:", two_sided_p(32856))
print("片側検定でのp値:", one_sided_p(32856))
print("")
print("神取ゼミの総期数 S=35510")
print("両側検定でのp値:", two_sided_p(35510))
print("片側検定でのp値:", one_sided_p(35510))

"""
Explanation: <table>
    <thead>
        <tr>
            <th>

            </th>
            <th>
                シミュレーション
            </th>
            <th>
                正確分布
            </th>
            <th>
                近似分布
            </th>
        </tr>
    </thead>
    <tbody>
        <tr>
            <th>
                平均
            </th>
            <td>
                33332.48
            </td>
            <td>
                33333.33
            </td>
            <td>
                33333.33
            </td>
        </tr>
        <tr>
            <th>
                分散
            </th>
            <td>
                1081320.67
            </td>
            <td>
                1077777.78
            </td>
            <td>
                1077777.78
            </td>
        </tr>
        <tr>
            <th>
                下側1%
            </th>
            <td>
                30956.00
            </td>
            <td>
                30967.00
            </td>
            <td>
                30918.21
            </td>
        </tr>
        <tr>
            <th>
                下側2.5%
            </th>
            <td>
                31325.00
            </td>
            <td>
                31330.00
            </td>
            <td>
                31298.58
            </td>
        </tr>
        <tr>
            <th>
                下側5%
            </th>
            <td>
                31641.50
            </td>
            <td>
                31645.00
            </td>
            <td>
                31625.71
            </td>
        </tr>
        <tr>
            <th>
                下側10%
            </th>
            <td>
                32010.00
            </td>
            <td>
                32010.00
            </td>
            <td>
                32002.88
            </td>
        </tr>
        <tr>
            <th>
                中央値
            </th>
            <td>
                33324.00
            </td>
            <td>
                33322.00
            </td>
            <td>
                33333.33
            </td>
        </tr>
        <tr>
            <th>
                上側10%
            </th>
            <td>
                34671.00
            </td>
            <td>
                34671.00
            </td>
            <td>
                34663.79
            </td>
        </tr>
        <tr>
            <th>
                上側5%
            </th>
            <td>
                35063.00
            </td>
            <td>
                35059.00
            </td>
            <td>
                35040.96
            </td>
        </tr>
        <tr>
            <th>
                上側2.5%
            </th>
            <td>
                35406.00
            </td>
            <td>
                35399.00
            </td>
            <td>
                35368.09
            </td>
        </tr>
        <tr>
            <th>
                上側1%
            </th>
            <td>
                35802.50
            </td>
            <td>
                35797.00
            </td>
            <td>
                35748.46
            </td>
        </tr>
    </tbody>
</table>

となる。正規分布の方が負の二項分布よりもやや裾が厚いが、ほぼ同じ。
実験で使用されたセッションの平均期数を検定
尾山ゼミの実験で使用されたゲームの期数の平均は32.856期、神取ゼミの実験では35.51期であった。総期数はそれぞれ32856期, 35510期となる。 
総期数が負の二項分布に従っているとして、
 ①両側検定（S=33322 vs S $\neq$ 33322）
 ②片側検定（尾山ゼミ: S=33322 vs S < 33322, 神取ゼミ: S=33322 vs S > 33322）
の両方の下でのp値を求める（Negbin(1000, 0.03) の下で、観測値よりも極端な値の出る確率を求める）。
End of explanation
"""


