import matplotlib.pyplot as plt
import random 
from docx import Document
from docx.shared import Pt
import os                        

# 設定資料與圖片的資料夾路徑
folder_path = r"D:\Ethan\程式\python\bionmial distribution"
graph_folder_path = r"D:\Ethan\程式\python\bionmial distribution\binomial distribution graph"

# 用來儲存機率數據與拒絕域文字
Data = []
reject_region_text = []

# 階乘函式
def factorial(n:int):
    n=int(n)
    if n==1 or n==0:
        return 1
    return factorial(n-1)*n
f=factorial  

# 組合函式 C(n,m) = n! / (m! * (n - m)!)
def combination(n:int, m:int):
    n=int(n)
    m=int(m)
    if m > n :
        n,m = m,n
    return int(f(n)/(f(m)*f(n-m)))
c=combination  

# 二項分布公式 P(X=k) = c(n,k)*(p^k)*(q^(n-k))
def binomial_distribution(k:int, n:int, p:float):
    n=int(n)
    p=float(p)
    k=int(k)
    q=1-p
    return (c(n,k)*(p**k)*(q**(n-k)))
bd=binomial_distribution

# 繪製機率分布圖
def drawing(n:int, p:float, alpha:float):

    index = [i for i in range(n+1)]
    output=[bd(i,n,p) for i in range(n+1)]
    Data.append({f'B({n},{p:.3f})': 
    [f'P(X={i})={data}' for i, data in enumerate(output)]})

    filename = "B({},{}).png".format(n,f"{p:.3f}")    
    
    # 繪圖設定
    fig, ax = plt.subplots()
    x_labels = [f"k={i}" for i in index]
    bars = ax.bar(x_labels, output)
    ax.set_title(f'B({n},{p})')
    ax.set_ylabel("P")
    ax.set_xlabel("X=k")

    # 找出拒絕域
    i=0
    while True:
        if sum(output[i::-1]+output[-1*i-1::1])>= alpha:
            i-=1
            reject_region = index[i::-1]+index[-1*i-1::1]
            reject_region_text.append(
            f"X={', '.join(str(k) for k in sorted(reject_region))}是拒絕域"
            )
            break
        i+=1
    
    # 將拒絕域標成紅色
    for bar in bars:
        for num in output[i::-1]+output[-1*i-1::1]:
            if bar.get_height()==num:
                bar.set_color("red")    

    for label in ax.get_xticklabels():
        for index in reject_region:
            if label.get_text() == f"k={index}":
                label.set_color("red")

    # 儲存圖表為PNG
    plt.savefig(os.path.join(graph_folder_path, filename))
    plt.close()


# 刪除舊的 PNG 圖片檔
for file in os.listdir(graph_folder_path):
    if os.path.exists(os.path.join(graph_folder_path, file)):
        os.remove(os.path.join(graph_folder_path, file))

# 刪除舊的 Word 文件
if os.path.exists(os.path.join(folder_path,'binomial distribution.docx')):
    os.remove(os.path.join(folder_path,'binomial distribution.docx'))

# 生成十組隨機 p 值的B(100,p) 顯著水準為0.05
for _ in range(5):
    drawing(200,random.random(),0.05)

# 建立 Word 文件並插入資料與圖片
document = Document()

document.add_heading('Binomial Distribution', 0)
for data in Data:
        for key, val in data.items():
            document.add_paragraph(key)  
        for text in reject_region_text:
            document.add_paragraph(text).style.font.size = Pt(16)
            for p in val:
                document.add_paragraph(p)
            document.add_picture(os.path.join(graph_folder_path,key+".png"))

document.save('binomial distribution.docx')