"""
Program obliczający frekwencje alleli w kolejnych pokoleniach, 
w zależności od ich początkowych frekwencji, oraz 
wartości dostosowania posczególnych genotypów.
"""

# Wprowadzenie początkowych wartości
q = float(input("podaj początkową frekwencję allelu a: "))
w_AA = float(input("podaj dostosowanie genotypu AA:    "))
w_Aa = float(input("podaj dostosowanie genotypu Aa:    "))
w_aa = float(input("podaj dostosowanie genotypu aa:    "))
max_pokolen = int(input("podaj maksymalną liczbę pokoleń:   "))

# Wartość p
p = 1 - q
# Pokolenie
n = 0
# Frekwencje początkowe
# Obliczanie frekwencji genotypów wg.
# prawa Hardy'ego-Weinberga
P_AA = p ** 2
P_Aa = 2 * p * q
P_aa = q ** 2
# Drukowanie Frekwencji początkowych:
print(f"F. początkowe - p_A:{p:.10f}, p_a: {q:.10f}, P_AA: {P_AA:.10f}, \
P_Aa: {P_Aa:.10f}, P_aa: {P_aa:.10f}")   
# Różnica frekwencji miedzy pokoleniami, ustalamy fikcyjną
# wartość, żeby pętla się uruchomiła
dq = 1
# Pętla wyliczająca frekwencje
while q > 0 and q < 1 and n <= max_pokolen and dq != 0:
    # Obliczanie frekwencji allelu a 
    # w kolejnym pokoleniu
    # średnie dostosowanie
    Ws = p ** 2 * w_AA + 2 * p * q * w_Aa + q ** 2 * w_aa
    # Zmiana frekwencji allelu a
    dq = (p * q * ( q * (w_aa - w_Aa) - p * (w_AA - w_Aa))) / Ws
    # Frekwencja allelu a w kolejnym pokoleniu
    q = q + dq
    # następne pokolenie
    n += 1
    # Obliczanie frekwencji allelu A
    p = 1 - q
    # Obliczanie frekwencji genotypów wg.
    # prawa Hardy'ego-Weinberga
    P_AA = p ** 2
    P_Aa = 2 * p * q
    P_aa = q ** 2
    # Drukowanie wartości
    #print(f" pok.: {n}, p_A:{p:.3f}, p_a: {q:.3f}, P_AA: {P_AA:.3f}, P_Aa: {P_Aa:.3f}, P_aa: {P_aa:.3f}")
    print(f"pok.: {n}, p_A:{p:.10f}, p_a: {q:.10f}, P_AA: {P_AA:.10f}, \
P_Aa: {P_Aa:.10f}, P_aa: {P_aa:.10f}, dq: {dq}")