import pandas as pd
import math
import matplotlib.pyplot as plt
# disable chained assignments
pd.options.mode.chained_assignment = None


def draw_F(x, p, show):
    print("Найдем функция распределения случайной величины.\nРисуем график F(x) от x")
    slv = "F(x)=\n"
    slv += "----\n"
    slv += f"|0 ,если x <={x[0]}\n"
    plt.plot([x[0]-2, x[0]], [0, 0], label=f"x <={x[0]}")
    for i in range(len(x)-1):
        sum_p = sum(p[0:i+1])
        plt.plot([x[i], x[i+1]], [sum_p, sum_p], label=f"{x[i]}< x <={x[i+1]}")
        slv += f"|{sum_p} ,если {x[i]}< x <={x[i+1]}\n"
    slv += f"|1 , если x >{x[len(x)-1]}\n"
    plt.plot([x[len(x)-1], x[len(x)-1]+2], [1, 1],  label=f"x <={x[0]}")
    slv += "----"
    if(show):
        print(slv)
    print([x[len(x)-1]])
    plt.xlabel('x')
    plt.ylabel('F(x)')
    plt.grid()
    plt.legend()
    if(show):
        plt.show()


def math_wait(x: list, p: list, show):
    x_graph = x
    p_graph = p
    slv = "Найдем Мат. Ожидание\n"
    slv += "M(X)= Σm(i)p(i) = "
    M = sum(float(x[i])*float(p[i]) for i in range(len(x)))
    slv += f"{M}"
    print(slv)

    slv = "Найдем \nx - M (x),\n(x - M (x))^2,\nx_i*p_i\nx_i^2*p_i,\nДисперсию = D(X)=M[(X - M(x^2))], которая не может быть отрицателной,\nсигму (называют средним квадратическим отклонением) σ(X)=√D(x)"
    print(slv)

    x_minus_Mx = [i-M for i in x]
    x_minus_Mx2 = [i**2 for i in x_minus_Mx]
    x_minus_Mx2_pi = [x_minus_Mx2[i]*p[i] for i in range(len(x_minus_Mx2))]
    xp = [x[i]*p[i] for i in range(len(x))]
    x2p = [(x[i]**2)*p[i] for i in range(len(x))]

    slv = "\nНайдем Дисперпсию cлучайной величины X\nD(X)=M[(X - M(x^2))] =\nили\nD(X)=M(X^2)-(M(X))^2=\n= "
    D = round(sum(x_minus_Mx2_pi), 3)
    print(pd.Series([sum(x)]))
    x = x.append(pd.Series([sum(x)]), ignore_index=True)
    p = p.append(pd.Series([sum(p)]), ignore_index=True)
    x_minus_Mx.append(sum(x_minus_Mx))
    x_minus_Mx2.append(sum(x_minus_Mx2))
    x_minus_Mx2_pi.append(sum(x_minus_Mx2_pi))
    xp.append(sum(xp))
    x2p.append(sum(x2p))
    # добавляем последний столбец - сумму
    df = pd.DataFrame({
        "x": x,
        "p": p,
        "x - M (x)": x_minus_Mx,
        "(x - M (x))^2": x_minus_Mx2,
        "(x-M(x))^2*p_i, sum=D": x_minus_Mx2_pi,
        "x_i*p_i": xp,
        "(x_i^2)*p_i": x2p,
    })
    df = df.transpose()
    df.rename(columns={len(p)-1: 'Сумма'}, inplace=True)
    print("\nТаблица:")
    print(df)
    if(D < 0):
        print("Ошибка! Дисперсия не может быть отрицательной")
    slv += f"{D}\n\n"
    sigma = round(math.sqrt(D), 3)
    slv += f"Найдем среднее квадратическое отклонение (сигма σ)\nσ(X)=√D(x) =√({D}) = "
    slv += f"{sigma}\n"
    print(slv)
    print(f"Ответ: M(X) = {M}, D(X) = {D}, σ(X) = {sigma}")

    print("Рисуем график pi от xi")
    plt.plot(x_graph, p_graph, color="blue", label="pi от xi")
    plt.xlabel('xi')
    plt.ylabel('pi')
    plt.grid()
    plt.legend()
    plt.show()
    draw_F(x_graph, p_graph, show)
    return M


def solve_table_p(table, show):
    print("----------")
    print(table)
    # решает таблицу и находит одну отсутвующую p
    table = list(map(list, zip(*table)))
    df = pd.DataFrame(table, columns=['m', 'p'])
    sum_p = 0
    p_index = -1
    for i in range(len(df["p"])):
        if type(df["p"][i]) != str:
            sum_p += df["p"][i]
        else:
            p_index = i
    if p_index == -1:
        return math_wait(df["m"], df["p"], show)
    x = round(1 - sum_p, 3)
    df["p"][p_index] = x
    print(f"Найдено отсутсвующее значение для {df['m'][p_index]} = {x}")
    df = df.sort_values(by=['m'])
    df2 = df.transpose()
    print("Вся таблица")
    print(df2)
    math_wait(df["m"], df["p"], show)
    return x


def corMoment(M_e, n_list: list, M_n, e_list: list, table):
    solve = "Корреляционный момент\n"
    solve += "K_xy =  ∑_i∑_j(x_i-m_x(y_i-m_i))pij = "
    K = 0
    for i in range(len(n_list)):
        for j in range(len(e_list)):
            K += (e_list[j]-M_e)*(n_list[i]-M_n)*table[j][i]
    solve += str(K)
    print(solve)
    print("Такие случайные величины называются некоррелированными (независимыми).")


def countQ(n_list, e_list, table, a, b):
    def q(e, n):
        return a*e+b*n
    row_of_distribution = []
    solve = ""
    for i in range(len(e_list)):
        for t in range(len(n_list)):
            q_buff = q(e_list[i], n_list[t])
            solve += f"Q_{i+1}{t+1} = {q_buff}\n"
            solve += f"P_{i+1}{t+1} = {table[i][t]}\n"
            row_of_distribution.append([q_buff, table[i][t]])
    row_of_distribution.sort(key=lambda x: x[0])
    print(solve)
    for row in row_of_distribution:
        print(row)
    solve_table_p([list(i) for i in zip(*row_of_distribution)], True)


def sumTable(n_list, e_list, table, a=0, b=0):
    solve = ""
    # e table
    solve += "Ei | "
    for i in range(len(e_list)):
        solve += str(e_list[i]) + "  |  "
    solve += "\n------------------------\n"
    solve += "Pi | "
    e_sum = []
    for i in range(len(e_list)):
        buff_sum = sum(table[i])
        e_sum.append(buff_sum)
        solve += str(round(buff_sum, 2)) + " | "

    solve += "\n\n"
    # n table
    solve += "Ni | "
    for i in range(len(n_list)):
        solve += str(n_list[i]) + "  |  "
    solve += "\n------------------\n"
    solve += "Pi | "
    n_sum = []
    for i in range(len(n_list)):
        buff_sum = sum([*zip(*table)][i])
        n_sum.append(buff_sum)
        solve += str(round(buff_sum, 2)) + " | "
    solve += "\n\n"
    print(solve)
    print("Найдем мат ожидание для e")
    M_e = solve_table_p([e_list, e_sum], True)
    M_n = solve_table_p([n_list, n_sum], True)
    corMoment(M_e, n_list, M_n, e_list, table)
    # if(a!=0):
    #     countQ(n_list, e_list, table, a, b)

n = [-1, 4]
e = [-1, 2, 3]
table = [
    [0.5, 0.1],
    [0.2, 0.05],
    [0.05, 0.1],
]
# вертикально записываем всегда большее количество чисел
if len(n) > len(e):
    # переворачиваем лист
    table = [*zip(*table)]
sumTable(n, e, table)