import numpy as np
import matplotlib.pyplot as plt

# Hardy-Weinberg Equilibrium formula:
# p^2 + 2pq + q^2 = 1
# where p is the frequency of allele A and q is the frequency of allele a

def compute_hwe(p, q):
    # Ensure that p and q add up to 1
    if p + q != 1:
        raise ValueError("The sum of allele frequencies p and q must be 1.")
    
    # Calculating the genotype frequencies
    p2 = p**2
    pq = 2 * p * q
    q2 = q**2
    
    return p2, pq, q2

# Range of allele frequencies
p_values = np.linspace(0, 1, 100)
q_values = 1 - p_values

# Store genotype frequencies
p2_values = []
pq_values = []
q2_values = []

# Calculate genotype frequencies for each allele frequency
for p, q in zip(p_values, q_values):
    p2, pq, q2 = compute_hwe(p, q)
    p2_values.append(p2)
    pq_values.append(pq)
    q2_values.append(q2)

# Convert lists to arrays
p2_values = np.array(p2_values)
pq_values = np.array(pq_values)
q2_values = np.array(q2_values)

# Plot the results
plt.figure(figsize=(10, 6))
plt.plot(p_values, p2_values, label='$p^2$ (AA)', color='blue')
plt.plot(p_values, pq_values, label='$2pq$ (Aa)', color='green')
plt.plot(p_values, q2_values, label='$q^2$ (aa)', color='red')

plt.xlabel('Allele Frequency $p$')
plt.ylabel('Genotype Frequency')
plt.title('Hardy-Weinberg Equilibrium')
plt.legend()
plt.grid(True)
plt.show()

# Brief description of computational genetics
"""
# Computational Genetics

Computational genetics is a field of study that applies computational and statistical 
methods to analyze and interpret genetic data. This interdisciplinary domain combines 
principles from genetics, computer science, mathematics, and statistics to understand 
the structure, function, and evolution of genomes.

One of the key concepts in computational genetics is the Hardy-Weinberg Equilibrium (HWE), 
which provides a mathematical model to study allele and genotype frequencies in a population. 
The HWE principle states that allele and genotype frequencies will remain constant from 
generation to generation in an ideal population that is not affected by mutation, selection,
migration, or genetic drift. This equilibrium can be described by the equation p^2 + 2pq + q^2 = 1, 
where p and q represent the frequencies of two alleles of a gene.

Computational genetics involves the use of algorithms and software tools to 
analyze large datasets generated by genomic sequencing technologies. These analyses 
can reveal important insights into genetic variation, gene function, and the genetic
basis of diseases. Researchers in this field may develop and apply methods for sequence
alignment, genome assembly, variant calling, and population genetics studies.

By integrating genetic data with computational techniques, scientists can uncover
patterns and relationships that would be difficult to detect using traditional methods.
This approach enables the identification of genetic markers associated with diseases, 
the study of evolutionary processes, and the development of personalized medicine. 
Computational genetics is a rapidly advancing field that continues to expand our 
understanding of the genetic basis of life and disease.
"""
