import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D

def naca4_airfoil(m, p, t, num_points=100):
    """
    Generate coordinates for a NACA 4-digit airfoil.
    
    Parameters:
    - m: maximum camber as fraction of chord
    - p: location of maximum camber as fraction of chord
    - t: maximum thickness as fraction of chord
    - num_points: number of points to generate
    
    Returns:
    - x_upper, y_upper: coordinates of the upper surface
    - x_lower, y_lower: coordinates of the lower surface
    """
    # Generate x-coordinates
    beta = np.linspace(0, np.pi, num_points)
    x = 0.5 * (1 - np.cos(beta))  # Cosine spacing for better resolution near leading edge
    
    # Calculate thickness distribution
    y_t = 5 * t * (0.2969 * np.sqrt(x) - 0.1260 * x - 0.3516 * x**2 + 0.2843 * x**3 - 0.1015 * x**4)
    
    if m == 0:
        # Symmetric airfoil
        x_upper = x
        y_upper = y_t
        x_lower = x
        y_lower = -y_t
    else:
        # Calculate camber line and its derivative
        y_c = np.zeros_like(x)
        dyc_dx = np.zeros_like(x)
        
        # For x < p
        mask = x <= p
        y_c[mask] = m * (x[mask] / p**2) * (2 * p - x[mask])
        dyc_dx[mask] = 2 * m / p**2 * (p - x[mask])
        
        # For x >= p
        mask = x > p
        y_c[mask] = m * ((1 - x[mask]) / (1 - p)**2) * (1 + x[mask] - 2 * p)
        dyc_dx[mask] = 2 * m / (1 - p)**2 * (p - x[mask])
        
        # Calculate theta
        theta = np.arctan(dyc_dx)
        
        # Calculate upper and lower surface coordinates
        x_upper = x - y_t * np.sin(theta)
        y_upper = y_c + y_t * np.cos(theta)
        x_lower = x + y_t * np.sin(theta)
        y_lower = y_c - y_t * np.cos(theta)
    
    return x_upper, y_upper, x_lower, y_lower

def plot_airfoil(x_upper, y_upper, x_lower, y_lower, title):
    """Plot an airfoil profile."""
    plt.figure(figsize=(10, 6))
    plt.plot(x_upper, y_upper, 'b-', label='Upper Surface')
    plt.plot(x_lower, y_lower, 'r-', label='Lower Surface')
    plt.grid(True)
    plt.axis('equal')
    plt.xlabel('x/c')
    plt.ylabel('y/c')
    plt.title(title)
    plt.legend()
    plt.show()

def calculate_lift_coefficient(alpha, camber, max_camber_pos):
    """
    Calculate approximate lift coefficient based on thin airfoil theory.
    
    Parameters:
    - alpha: angle of attack in degrees
    - camber: maximum camber
    - max_camber_pos: position of maximum camber
    
    Returns:
    - CL: lift coefficient
    """
    alpha_rad = np.deg2rad(alpha)
    
    # Simplified thin airfoil theory approximation
    CL = 2 * np.pi * alpha_rad
    
    # Add contribution from camber
    if camber > 0:
        # Simplified approximation for camber contribution
        CL += np.pi * camber * (2 - 4 * max_camber_pos)
    
    return CL

def plot_lift_curve(airfoil_name, m, p, t, alpha_range):
    """
    Plot lift curve for a given airfoil.
    
    Parameters:
    - airfoil_name: name for plot title
    - m, p, t: airfoil parameters
    - alpha_range: range of angles of attack to plot
    """
    lift_coefficients = [calculate_lift_coefficient(alpha, m, p) for alpha in alpha_range]
    
    plt.figure(figsize=(10, 6))
    plt.plot(alpha_range, lift_coefficients, 'b-o')
    plt.grid(True)
    plt.xlabel('Angle of Attack (degrees)')
    plt.ylabel('Lift Coefficient (CL)')
    plt.title(f'Lift Curve for {airfoil_name}')
    plt.axhline(y=0, color='k', linestyle='-', alpha=0.3)
    plt.axvline(x=0, color='k', linestyle='-', alpha=0.3)
    plt.show()
    
    return lift_coefficients

def study_camber_effect():
    """Study the effect of camber on airfoil performance."""
    # Define a range of camber values
    camber_values = np.linspace(0, 0.06, 4)  # 0% to 6% camber
    alpha = 5  # Fixed angle of attack (5 degrees)
    max_camber_pos = 0.4  # Fixed position of maximum camber
    thickness = 0.12  # Fixed thickness
    
    plt.figure(figsize=(12, 8))
    
    # For each camber value, generate and plot the airfoil
    for m in camber_values:
        x_upper, y_upper, x_lower, y_lower = naca4_airfoil(m, max_camber_pos, thickness)
        label = f'Camber = {m:.2f}'
        plt.plot(x_upper, y_upper, '-', label=label)
        plt.plot(x_lower, y_lower, '-')
    
    plt.grid(True)
    plt.axis('equal')
    plt.xlabel('x/c')
    plt.ylabel('y/c')
    plt.title('Effect of Camber on Airfoil Shape')
    plt.legend()
    plt.show()
    
    # Calculate and plot lift coefficient vs camber
    camber_range = np.linspace(0, 0.1, 50)
    lift_coefficients = [calculate_lift_coefficient(alpha, m, max_camber_pos) for m in camber_range]
    
    plt.figure(figsize=(10, 6))
    plt.plot(camber_range, lift_coefficients, 'b-')
    plt.grid(True)
    plt.xlabel('Maximum Camber (m)')
    plt.ylabel('Lift Coefficient (CL) at α = 5°')
    plt.title('Effect of Camber on Lift Coefficient')
    plt.show()

def study_thickness_effect():
    """Study the effect of thickness on airfoil shape."""
    thickness_values = np.linspace(0.06, 0.18, 4)  # 6% to 18% thickness
    camber = 0.04  # Fixed camber
    max_camber_pos = 0.4  # Fixed position of maximum camber
    
    plt.figure(figsize=(12, 8))
    
    # For each thickness value, generate and plot the airfoil
    for t in thickness_values:
        x_upper, y_upper, x_lower, y_lower = naca4_airfoil(camber, max_camber_pos, t)
        label = f'Thickness = {t:.2f}'
        plt.plot(x_upper, y_upper, '-', label=label)
        plt.plot(x_lower, y_lower, '-')
    
    plt.grid(True)
    plt.axis('equal')
    plt.xlabel('x/c')
    plt.ylabel('y/c')
    plt.title('Effect of Thickness on Airfoil Shape')
    plt.legend()
    plt.show()

def pressure_distribution(x_upper, y_upper, x_lower, y_lower, alpha=5):
    """
    Calculate approximate pressure coefficient distribution around an airfoil.
    
    Parameters:
    - x_upper, y_upper, x_lower, y_lower: airfoil coordinates
    - alpha: angle of attack in degrees
    
    Returns:
    - Cp_upper, Cp_lower: pressure coefficients for upper and lower surfaces
    """
    # Convert angle of attack to radians
    alpha_rad = np.deg2rad(alpha)
    
    # Calculate local velocity (simplified model based on thin airfoil theory)
    # For upper surface
    Cp_upper = 1 - (2 * np.sin(alpha_rad + np.arctan2(np.gradient(y_upper, x_upper), 1)))**2
    
    # For lower surface
    Cp_lower = 1 - (2 * np.sin(alpha_rad + np.arctan2(np.gradient(y_lower, x_lower), 1)))**2
    
    return Cp_upper, Cp_lower

def plot_pressure_distribution(x_upper, y_upper, x_lower, y_lower, alpha, title):
    """Plot pressure coefficient distribution around an airfoil."""
    Cp_upper, Cp_lower = pressure_distribution(x_upper, y_upper, x_lower, y_lower, alpha)
    
    plt.figure(figsize=(12, 6))
    plt.plot(x_upper, -Cp_upper, 'b-', label='Upper Surface')
    plt.plot(x_lower, -Cp_lower, 'r-', label='Lower Surface')
    plt.grid(True)
    plt.xlabel('x/c')
    plt.ylabel('-Cp')
    plt.title(f'Pressure Distribution around {title} at α = {alpha}°')
    plt.legend()
    plt.gca().invert_yaxis()  # Invert y-axis to match convention (suction at top)
    plt.show()

def plot_3d_pressure_flow(airfoil_name, m, p, t, alpha_range):
    """Create a 3D visualization of pressure distribution at different angles of attack."""
    # Generate airfoil
    x_upper, y_upper, x_lower, y_lower = naca4_airfoil(m, p, t, num_points=50)
    
    # Prepare 3D plot data
    X, Y = np.meshgrid(x_upper, alpha_range)
    Z = np.zeros_like(X)
    
    # Calculate pressure distribution for each angle of attack
    for i, alpha in enumerate(alpha_range):
        Cp_upper, _ = pressure_distribution(x_upper, y_upper, x_lower, y_lower, alpha)
        Z[i, :] = -Cp_upper  # Negative Cp to show suction as positive
    
    # Create 3D plot
    fig = plt.figure(figsize=(12, 10))
    ax = fig.add_subplot(111, projection='3d')
    
    # Plot surface
    surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm, linewidth=0, antialiased=True)
    
    # Add labels and colorbar
    ax.set_xlabel('x/c')
    ax.set_ylabel('Angle of Attack (degrees)')
    ax.set_zlabel('-Cp (Pressure Coefficient)')
    ax.set_title(f'3D Pressure Distribution for {airfoil_name}')
    fig.colorbar(surf, shrink=0.5, aspect=5)
    
    plt.show()

# Main execution
if __name__ == "__main__":
    print("NACA Airfoil Analysis Example\n")
    
    # Define example airfoil parameters
    # NACA 4412 Airfoil (4% camber, 40% max camber position, 12% thickness)
    m = 0.04  # Maximum camber
    p = 0.4   # Position of maximum camber
    t = 0.12  # Maximum thickness
    
    airfoil_name = f"NACA 4412"
    print(f"Analyzing {airfoil_name} airfoil:")
    print(f"- Maximum camber: {m*100}%")
    print(f"- Position of maximum camber: {p*100}%")
    print(f"- Maximum thickness: {t*100}%\n")
    
    # Generate airfoil coordinates
    x_upper, y_upper, x_lower, y_lower = naca4_airfoil(m, p, t)
    
    # Plot airfoil profile
    plot_airfoil(x_upper, y_upper, x_lower, y_lower, airfoil_name)
    
    # Calculate and plot lift curve
    alpha_range = np.linspace(-5, 15, 21)
    lift_coefficients = plot_lift_curve(airfoil_name, m, p, t, alpha_range)
    
    # Plot pressure distribution at 5 degrees angle of attack
    plot_pressure_distribution(x_upper, y_upper, x_lower, y_lower, 5, airfoil_name)
    
    # Study effect of camber on airfoil performance
    study_camber_effect()
    
    # Study effect of thickness on airfoil shape
    study_thickness_effect()
    
    # Create 3D visualization of pressure distribution
    plot_3d_pressure_flow(airfoil_name, m, p, t, np.linspace(0, 10, 11))
    
    print("Analysis complete!")

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from scipy.stats import chi2_contingency
import seaborn as sns

# Set seed for reproducibility
np.random.seed(42)

def simulate_cross(genotype1, genotype2, num_offspring=1000):
    """
    Simulate a genetic cross between two genotypes.
    
    Parameters:
    genotype1, genotype2: Strings representing parental genotypes (e.g., 'Aa', 'BB')
    num_offspring: Number of offspring to simulate
    
    Returns:
    Dictionary with genotype counts and phenotype counts
    """
    # Extract alleles from genotypes
    alleles1 = [genotype1[0], genotype1[1]]
    alleles2 = [genotype2[0], genotype2[1]]
    
    # Generate gametes
    gametes1 = np.random.choice(alleles1, size=num_offspring)
    gametes2 = np.random.choice(alleles2, size=num_offspring)
    
    # Combine gametes to form offspring genotypes
    offspring_genotypes = np.array([g1 + g2 for g1, g2 in zip(gametes1, gametes2)])
    
    # Count genotypes
    unique_genotypes, genotype_counts = np.unique(offspring_genotypes, return_counts=True)
    genotype_freqs = {gt: count/num_offspring for gt, count in zip(unique_genotypes, genotype_counts)}
    
    # Determine phenotypes (assuming dominant/recessive relationship)
    # If capital letter is present, express dominant phenotype
    phenotypes = []
    for gt in offspring_genotypes:
        if gt[0].isupper() or gt[1].isupper():
            phenotypes.append("Dominant")
        else:
            phenotypes.append("Recessive")
    
    unique_phenotypes, phenotype_counts = np.unique(phenotypes, return_counts=True)
    phenotype_freqs = {pt: count/num_offspring for pt, count in zip(unique_phenotypes, phenotype_counts)}
    
    return {
        "genotype_counts": {gt: count for gt, count in zip(unique_genotypes, genotype_counts)},
        "genotype_freqs": genotype_freqs,
        "phenotype_counts": {pt: count for pt, count in zip(unique_phenotypes, phenotype_counts)},
        "phenotype_freqs": phenotype_freqs,
        "raw_genotypes": offspring_genotypes,
        "raw_phenotypes": phenotypes
    }

# Example 1: Single gene inheritance (Aa × Aa cross)
print("Simulating monohybrid cross: Aa × Aa")
monohybrid_results = simulate_cross("Aa", "Aa", 1000)

print("\nGenotype counts:")
for genotype, count in monohybrid_results["genotype_counts"].items():
    print(f"{genotype}: {count} ({count/10:.1f}%)")

print("\nPhenotype counts:")
for phenotype, count in monohybrid_results["phenotype_counts"].items():
    print(f"{phenotype}: {count} ({count/10:.1f}%)")

# Chi-square test to check if our results match Mendelian ratios
expected_genotype_ratio = {"AA": 0.25, "Aa": 0.5, "aA": 0, "aa": 0.25}
# Adjust for the fact that Aa and aA are equivalent
if "Aa" in monohybrid_results["genotype_counts"] and "aA" in monohybrid_results["genotype_counts"]:
    expected_genotype_ratio["Aa"] = 0.5
elif "aA" in monohybrid_results["genotype_counts"]:
    expected_genotype_ratio["aA"] = 0.5

observed = []
expected = []
for genotype in monohybrid_results["genotype_counts"]:
    if genotype in expected_genotype_ratio:
        observed.append(monohybrid_results["genotype_counts"][genotype])
        expected.append(expected_genotype_ratio[genotype] * 1000)

chi2, p, dof, ex = chi2_contingency([observed, expected])
print(f"\nChi-square test p-value: {p:.4f}")
if p > 0.05:
    print("The observed results are consistent with Mendelian inheritance.")
else:
    print("The observed results differ significantly from Mendelian inheritance.")

# Example 2: Dihybrid cross (AaBb × AaBb)
def simulate_dihybrid_cross(genotype1, genotype2, num_offspring=1000):
    """Simulate a dihybrid cross between two genotypes"""
    # Split genotypes into separate genes
    gene1_p1, gene2_p1 = genotype1[:2], genotype1[2:]
    gene1_p2, gene2_p2 = genotype2[:2], genotype2[2:]
    
    # Simulate each gene independently
    gene1_results = simulate_cross(gene1_p1, gene1_p2, num_offspring)
    gene2_results = simulate_cross(gene2_p1, gene2_p2, num_offspring)
    
    # Combine results to get dihybrid genotypes
    dihybrid_genotypes = np.array([g1 + g2 for g1, g2 in 
                                 zip(gene1_results["raw_genotypes"], gene2_results["raw_genotypes"])])
    
    # Count genotypes
    unique_genotypes, genotype_counts = np.unique(dihybrid_genotypes, return_counts=True)
    genotype_freqs = {gt: count/num_offspring for gt, count in zip(unique_genotypes, genotype_counts)}
    
    # Determine phenotypes
    phenotypes = []
    for gt in dihybrid_genotypes:
        p1 = "A_dominant" if (gt[0].isupper() or gt[1].isupper()) else "a_recessive"
        p2 = "B_dominant" if (gt[2].isupper() or gt[3].isupper()) else "b_recessive"
        phenotypes.append(f"{p1}_{p2}")
    
    unique_phenotypes, phenotype_counts = np.unique(phenotypes, return_counts=True)
    phenotype_freqs = {pt: count/num_offspring for pt, count in zip(unique_phenotypes, phenotype_counts)}
    
    # Simplify phenotype names for clarity in plotting
    simplified_phenotypes = []
    for p in phenotypes:
        if p == "A_dominant_B_dominant":
            simplified_phenotypes.append("A-B")
        elif p == "A_dominant_b_recessive":
            simplified_phenotypes.append("A-b")
        elif p == "a_recessive_B_dominant":
            simplified_phenotypes.append("a-B")
        elif p == "a_recessive_b_recessive":
            simplified_phenotypes.append("a-b")
    
    # Count simplified phenotypes
    unique_simple_phenotypes, simple_phenotype_counts = np.unique(simplified_phenotypes, return_counts=True)
    simple_phenotype_freqs = {pt: count/num_offspring for pt, count in zip(unique_simple_phenotypes, simple_phenotype_counts)}
    
    return {
        "genotype_counts": {gt: count for gt, count in zip(unique_genotypes, genotype_counts)},
        "genotype_freqs": genotype_freqs,
        "phenotype_counts": {pt: count for pt, count in zip(unique_phenotypes, phenotype_counts)},
        "phenotype_freqs": phenotype_freqs,
        "simplified_phenotype_counts": {pt: count for pt, count in zip(unique_simple_phenotypes, simple_phenotype_counts)},
        "simplified_phenotype_freqs": simple_phenotype_freqs
    }

print("\n\nSimulating dihybrid cross: AaBb × AaBb")
dihybrid_results = simulate_dihybrid_cross("AaBb", "AaBb", 1000)

# Plot the results
plt.figure(figsize=(14, 10))

# Plot 1: Monohybrid cross genotype frequencies
plt.subplot(2, 2, 1)
genotypes = list(monohybrid_results["genotype_counts"].keys())
counts = list(monohybrid_results["genotype_counts"].values())
plt.bar(genotypes, counts, color='skyblue')
plt.title('Monohybrid Cross (Aa × Aa)\nGenotype Frequencies')
plt.ylabel('Count')
for i, v in enumerate(counts):
    plt.text(i, v + 5, f"{v}", ha='center')

# Plot 2: Monohybrid cross phenotype frequencies
plt.subplot(2, 2, 2)
phenotypes = list(monohybrid_results["phenotype_counts"].keys())
counts = list(monohybrid_results["phenotype_counts"].values())
plt.bar(phenotypes, counts, color='lightgreen')
plt.title('Monohybrid Cross (Aa × Aa)\nPhenotype Frequencies')
plt.ylabel('Count')
for i, v in enumerate(counts):
    plt.text(i, v + 5, f"{v} ({v/10:.1f}%)", ha='center')

# Plot 3: Dihybrid cross simplified phenotype frequencies as pie chart
plt.subplot(2, 2, 3)
labels = list(dihybrid_results["simplified_phenotype_counts"].keys())
sizes = list(dihybrid_results["simplified_phenotype_counts"].values())
explode = (0.1, 0, 0, 0)  # explode the 1st slice
colors = ['#ff9999','#66b3ff','#99ff99','#ffcc99']
plt.pie(sizes, explode=explode, labels=labels, colors=colors, autopct='%1.1f%%', shadow=True, startangle=90)
plt.axis('equal')  # Equal aspect ratio ensures that pie is drawn as a circle
plt.title('Dihybrid Cross (AaBb × AaBb)\nPhenotype Ratio')

# Plot 4: Expected vs Observed for dihybrid cross (9:3:3:1 ratio)
plt.subplot(2, 2, 4)
phenotypes = list(dihybrid_results["simplified_phenotype_counts"].keys())
observed = list(dihybrid_results["simplified_phenotype_counts"].values())

# Expected values based on 9:3:3:1 ratio
expected = [9/16*1000, 3/16*1000, 3/16*1000, 1/16*1000]

expected_labels = [f"{phenotypes[i]}\n(Exp: {expected[i]:.1f})" for i in range(len(phenotypes))]

x = np.arange(len(phenotypes))
width = 0.35

plt.bar(x - width/2, observed, width, label='Observed', color='royalblue')
plt.bar(x + width/2, expected, width, label='Expected (9:3:3:1)', color='lightcoral')
plt.title('Dihybrid Cross: Observed vs Expected')
plt.xticks(x, expected_labels, rotation=45, ha='right')
plt.ylabel('Count')
plt.legend()

plt.tight_layout()
plt.savefig('genetic_crosses.png', dpi=300)
plt.show()

# Additional Analysis: Linkage study
# Simulate crossing over and genetic linkage
def simulate_linkage(recombination_rate, num_offspring=1000):
    """
    Simulate genetic linkage between two genes with a given recombination rate.
    
    Parameters:
    recombination_rate: Probability of recombination between loci (0-0.5)
    num_offspring: Number of offspring to simulate
    
    Returns:
    Dictionary with gamete types and their frequencies
    """
    # Parent is double heterozygote in coupling phase (AB/ab)
    # Determine if recombination occurs for each offspring
    recombination = np.random.random(num_offspring) < recombination_rate
    
    # Generate gametes
    gametes = []
    for r in recombination:
        if r:  # Recombination occurred
            if np.random.random() < 0.5:
                gametes.append("Ab")  # Recombinant
            else:
                gametes.append("aB")  # Recombinant
        else:  # No recombination
            if np.random.random() < 0.5:
                gametes.append("AB")  # Parental
            else:
                gametes.append("ab")  # Parental
    
    # Count gamete types
    unique_gametes, gamete_counts = np.unique(gametes, return_counts=True)
    gamete_freqs = {gt: count/num_offspring for gt, count in zip(unique_gametes, gamete_counts)}
    
    return {
        "gamete_counts": {gt: count for gt, count in zip(unique_gametes, gamete_counts)},
        "gamete_freqs": gamete_freqs,
        "raw_gametes": gametes
    }

# Simulate different recombination rates
recom_rates = [0.01, 0.1, 0.2, 0.3, 0.4, 0.5]
linkage_results = []

for rate in recom_rates:
    result = simulate_linkage(rate, 1000)
    linkage_results.append(result)

# Plot linkage results
plt.figure(figsize=(14, 6))

# Plot: Effect of recombination rate on gamete frequencies
plt.subplot(1, 2, 1)
parental_freqs = []
recombinant_freqs = []

for i, rate in enumerate(recom_rates):
    result = linkage_results[i]
    
    # Sum parental gamete frequencies (AB + ab)
    parental = sum([result["gamete_counts"].get("AB", 0), 
                     result["gamete_counts"].get("ab", 0)]) / 1000
    
    # Sum recombinant gamete frequencies (Ab + aB)
    recombinant = sum([result["gamete_counts"].get("Ab", 0), 
                       result["gamete_counts"].get("aB", 0)]) / 1000
    
    parental_freqs.append(parental)
    recombinant_freqs.append(recombinant)

plt.plot(recom_rates, parental_freqs, 'o-', label='Parental Gametes (AB, ab)')
plt.plot(recom_rates, recombinant_freqs, 's-', label='Recombinant Gametes (Ab, aB)')
plt.axhline(y=0.5, color='gray', linestyle='--', alpha=0.5)
plt.title('Effect of Recombination Rate on Gamete Frequencies')
plt.xlabel('Recombination Rate')
plt.ylabel('Frequency')
plt.xticks(recom_rates)
plt.ylim(0, 1)
plt.legend()
plt.grid(alpha=0.3)

# Plot: Gamete distribution for different recombination rates
plt.subplot(1, 2, 2)
gamete_types = ["AB", "ab", "Ab", "aB"]
data = []

for i, rate in enumerate(recom_rates):
    result = linkage_results[i]
    row = [result["gamete_counts"].get(gt, 0)/1000 for gt in gamete_types]
    data.append(row)

df = pd.DataFrame(data, columns=gamete_types, index=[f"{rate*100:.0f}%" for rate in recom_rates])
sns.heatmap(df, annot=True, cmap="YlGnBu", fmt=".2f", cbar_kws={'label': 'Frequency'})
plt.title('Gamete Frequencies at Different Recombination Rates')
plt.xlabel('Gamete Type')
plt.ylabel('Recombination Rate')

plt.tight_layout()
plt.savefig('genetic_linkage.png', dpi=300)
plt.show()

# Example 3: Hardy-Weinberg Equilibrium
def simulate_hardy_weinberg(p_initial, num_generations=10, population_size=1000, selection_coefficient=0):
    """
    Simulate Hardy-Weinberg equilibrium over generations.
    
    Parameters:
    p_initial: Initial frequency of dominant allele A
    num_generations: Number of generations to simulate
    population_size: Size of the population in each generation
    selection_coefficient: Selection against the recessive homozygote (aa)
    
    Returns:
    Dictionary with allele and genotype frequencies over generations
    """
    q_initial = 1 - p_initial
    
    # Initialize arrays to store frequencies
    p_values = [p_initial]
    q_values = [q_initial]
    AA_freqs = [p_initial**2]
    Aa_freqs = [2 * p_initial * q_initial]
    aa_freqs = [q_initial**2]
    
    p = p_initial
    q = q_initial
    
    for gen in range(1, num_generations):
        # Calculate genotype frequencies under Hardy-Weinberg
        AA_freq = p**2
        Aa_freq = 2 * p * q
        aa_freq = q**2
        
        # Apply selection if specified
        if selection_coefficient > 0:
            # Relative fitness of genotypes
            w_AA = 1
            w_Aa = 1
            w_aa = 1 - selection_coefficient
            
            # Mean fitness
            w_bar = w_AA * AA_freq + w_Aa * Aa_freq + w_aa * aa_freq
            
            # New genotype frequencies after selection
            AA_freq = (w_AA * AA_freq) / w_bar
            Aa_freq = (w_Aa * Aa_freq) / w_bar
            aa_freq = (w_aa * aa_freq) / w_bar
            
            # New allele frequencies
            p = AA_freq + (Aa_freq / 2)
            q = aa_freq + (Aa_freq / 2)
        
        # Random sampling for finite population
        genotype_counts = np.random.multinomial(population_size, [AA_freq, Aa_freq, aa_freq])
        AA_count, Aa_count, aa_count = genotype_counts
        
        # Recalculate frequencies based on counts
        AA_freq = AA_count / population_size
        Aa_freq = Aa_count / population_size
        aa_freq = aa_count / population_size
        
        # Recalculate allele frequencies
        p = AA_freq + (Aa_freq / 2)
        q = aa_freq + (Aa_freq / 2)
        
        # Store frequencies
        p_values.append(p)
        q_values.append(q)
        AA_freqs.append(AA_freq)
        Aa_freqs.append(Aa_freq)
        aa_freqs.append(aa_freq)
    
    return {
        "p_values": p_values,
        "q_values": q_values,
        "AA_freqs": AA_freqs,
        "Aa_freqs": Aa_freqs,
        "aa_freqs": aa_freqs
    }

# Simulate Hardy-Weinberg with different initial conditions
hw_results_no_selection = simulate_hardy_weinberg(0.3, num_generations=20, population_size=1000, selection_coefficient=0)
hw_results_with_selection = simulate_hardy_weinberg(0.3, num_generations=20, population_size=1000, selection_coefficient=0.2)

# Plot Hardy-Weinberg results
plt.figure(figsize=(14, 10))

# Plot 1: Allele frequencies over generations (no selection)
plt.subplot(2, 2, 1)
generations = range(len(hw_results_no_selection["p_values"]))
plt.plot(generations, hw_results_no_selection["p_values"], 'b-', label='p (A allele)')
plt.plot(generations, hw_results_no_selection["q_values"], 'r-', label='q (a allele)')
plt.title('Allele Frequencies Over Generations\n(No Selection)')
plt.xlabel('Generation')
plt.ylabel('Frequency')
plt.legend()
plt.grid(alpha=0.3)

# Plot 2: Genotype frequencies over generations (no selection)
plt.subplot(2, 2, 2)
plt.plot(generations, hw_results_no_selection["AA_freqs"], 'g-', label='AA')
plt.plot(generations, hw_results_no_selection["Aa_freqs"], 'b-', label='Aa')
plt.plot(generations, hw_results_no_selection["aa_freqs"], 'r-', label='aa')
plt.title('Genotype Frequencies Over Generations\n(No Selection)')
plt.xlabel('Generation')
plt.ylabel('Frequency')
plt.legend()
plt.grid(alpha=0.3)

# Plot 3: Allele frequencies over generations (with selection)
plt.subplot(2, 2, 3)
plt.plot(generations, hw_results_with_selection["p_values"], 'b-', label='p (A allele)')
plt.plot(generations, hw_results_with_selection["q_values"], 'r-', label='q (a allele)')
plt.title('Allele Frequencies Over Generations\n(Selection Against aa, s=0.2)')
plt.xlabel('Generation')
plt.ylabel('Frequency')
plt.legend()
plt.grid(alpha=0.3)

# Plot 4: Genotype frequencies over generations (with selection)
plt.subplot(2, 2, 4)
plt.plot(generations, hw_results_with_selection["AA_freqs"], 'g-', label='AA')
plt.plot(generations, hw_results_with_selection["Aa_freqs"], 'b-', label='Aa')
plt.plot(generations, hw_results_with_selection["aa_freqs"], 'r-', label='aa')
plt.title('Genotype Frequencies Over Generations\n(Selection Against aa, s=0.2)')
plt.xlabel('Generation')
plt.ylabel('Frequency')
plt.legend()
plt.grid(alpha=0.3)

plt.tight_layout()
plt.savefig('hardy_weinberg.png', dpi=300)
plt.show()

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import linprog

# Problem: A company manufactures two products, A and B.
# Each product requires time on two machines: cutting and assembly.
# Machine constraints and profit information:
# - Product A: 2 hours cutting, 1 hour assembly, $40 profit per unit
# - Product B: 1 hour cutting, 3 hours assembly, $30 profit per unit
# - Cutting machine available 100 hours per week
# - Assembly machine available 90 hours per week
# Goal: Maximize profit by determining optimal production quantities

# Define the problem
# Objective function coefficients (negative because linprog minimizes)
c = [-40, -30]  # Profits for products A and B (negative for maximization)

# Inequality constraint coefficients (left side of <=)
A_ub = [
    [2, 1],    # Cutting machine hours for A and B
    [1, 3]     # Assembly machine hours for A and B
]

# Inequality constraint bounds (right side of <=)
b_ub = [100,  # Available cutting hours
        90]   # Available assembly hours

# Bounds for variables (non-negative)
x_bounds = [(0, None), (0, None)]  # Lower and upper bounds for A and B

# Solve the linear programming problem
result = linprog(c, A_ub, b_ub, bounds=x_bounds, method="highs")

# Print the results
print("Optimization successful:", result.success)
print("Optimal values:")
print(f"Product A: {result.x[0]:.2f} units")
print(f"Product B: {result.x[1]:.2f} units")
print(f"Maximum profit: ${-result.fun:.2f}")

# Visualization: Feasible region and optimal solution
fig, ax = plt.subplots(figsize=(10, 8))

# Create a grid of points
x1 = np.linspace(0, 60, 1000)

# Plot the constraints
# Cutting constraint: 2A + B <= 100 => B <= 100 - 2A
y1 = 100 - 2 * x1
# Assembly constraint: A + 3B <= 90 => B <= (90 - A) / 3
y2 = (90 - x1) / 3

# Plot the constraints
ax.plot(x1, y1, label="Cutting: 2A + B ≤ 100", linewidth=2)
ax.plot(x1, y2, label="Assembly: A + 3B ≤ 90", linewidth=2)

# Fill the feasible region
# Find the intersection of constraints
# 100 - 2x = (90 - x)/3
# 300 - 6x = 90 - x
# -5x = -210
# x = 42
# y = 100 - 2(42) = 16
feasible_x = [0, 0, result.x[0], 42, 50, 0]
feasible_y = [0, 30, result.x[1], 16, 0, 0]
ax.fill(feasible_x, feasible_y, alpha=0.3, color='green', label="Feasible Region")

# Plot the optimal solution
ax.scatter(result.x[0], result.x[1], color='red', s=100, 
           label=f"Optimal Solution: A={result.x[0]:.1f}, B={result.x[1]:.1f}")

# Set the axes limits and labels
ax.set_xlim(0, 60)
ax.set_ylim(0, 40)
ax.set_xlabel("Product A (units)", fontsize=12)
ax.set_ylabel("Product B (units)", fontsize=12)
ax.set_title("Resource Optimization Problem", fontsize=16)
ax.grid(True)
ax.legend(fontsize=10)

# Add isoprofit lines (objective function)
# Profit = 40A + 30B
# For various profit levels
profits = [1000, 1500, 2000, 2500]
colors = ['gray', 'lightgray', 'darkgray', 'black']
for i, profit in enumerate(profits):
    # Equation: 40A + 30B = profit => B = (profit - 40A) / 30
    y_profit = (profit - 40 * x1) / 30
    valid_indices = (y_profit >= 0) & (x1 >= 0)
    ax.plot(x1[valid_indices], y_profit[valid_indices], '--', color=colors[i], 
            label=f"Profit=${profit}")

# Recalculate the actual profit 
optimal_profit = 40 * result.x[0] + 30 * result.x[1]

# Add a text box with the optimal solution
textstr = '\n'.join((
    f'Optimal Solution:',
    f'Product A: {result.x[0]:.2f} units',
    f'Product B: {result.x[1]:.2f} units',
    f'Maximum Profit: ${optimal_profit:.2f}'))
props = dict(boxstyle='round', facecolor='wheat', alpha=0.5)
ax.text(1, 2, textstr, fontsize=12, verticalalignment='bottom', bbox=props)

plt.legend()
plt.tight_layout()
plt.show()

!pip install basemap
!pip install basemap-data-hires
!pip install networkx

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
import networkx as nx

def haversine(lon1, lat1, lon2, lat2):
    # Convert latitude and longitude to radians
    lon1, lat1, lon2, lat2 = map(np.radians, [lon1, lat1, lon2, lat2])

    # Haversine formula
    dlon = lon2 - lon1
    dlat = lat2 - lat1
    a = np.sin(dlat/2)**2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon/2)**2
    c = 2 * np.arcsin(np.sqrt(a))
    r = 6371  # Earth's radius in kilometers
    return c * r

# City names and their coordinates (latitude, longitude)
cities = {
    "New York": (40.7128, -74.0060),
    "London": (51.5074, -0.1278),
    "Tokyo": (35.6895, 139.6917),
    "Sydney": (-33.8688, 151.2093),
    "Cape Town": (-33.9249, 18.4241)
}

# Get the list of cities
city_names = list(cities.keys())
n_cities = len(city_names)

# Initialize the distance matrix
distance_matrix = np.zeros((n_cities, n_cities))

# Calculate the distance matrix
for i in range(n_cities):
    for j in range(n_cities):
        if i != j:
            lat1, lon1 = cities[city_names[i]]
            lat2, lon2 = cities[city_names[j]]
            distance_matrix[i, j] = haversine(lon1, lat1, lon2, lat2)
        else:
            distance_matrix[i, j] = 0  # Distance between the same city is 0

# Display the distance matrix
print("Distance Matrix (km):")
print(distance_matrix)

# Create a graph
G = nx.Graph()

# Add cities as nodes
for i, city in enumerate(city_names):
    G.add_node(i, name=city)

# Add edges (with distances as weights)
for i in range(n_cities):
    for j in range(i + 1, n_cities):
        G.add_edge(i, j, weight=distance_matrix[i, j])

# Compute the shortest path (e.g., from New York (0) to Tokyo (2))
start_city = 0  # New York
end_city = 2    # Tokyo
shortest_path = nx.dijkstra_path(G, start_city, end_city, weight='weight')
shortest_path_distance = nx.dijkstra_path_length(G, start_city, end_city, weight='weight')

# Display the shortest path
print(f"Shortest path from {city_names[start_city]} to {city_names[end_city]}:")
print(" -> ".join([city_names[i] for i in shortest_path]))
print(f"Total distance: {shortest_path_distance:.2f} km")

# Create a map
plt.figure(figsize=(12, 8))
m = Basemap(projection='merc', llcrnrlat=-60, urcrnrlat=70, llcrnrlon=-180, urcrnrlon=180, resolution='c')

# Draw coastlines, countries, and states
m.drawcoastlines()
m.drawcountries()
m.drawstates()

# Plot cities on the map
for i, (city, (lat, lon)) in enumerate(cities.items()):
    x, y = m(lon, lat)
    m.plot(x, y, 'ro', markersize=10)
    plt.text(x, y, city, fontsize=12, ha='right')

# Draw the shortest path
for i in range(len(shortest_path) - 1):
    start = shortest_path[i]
    end = shortest_path[i + 1]
    start_lon, start_lat = cities[city_names[start]][1], cities[city_names[start]][0]
    end_lon, end_lat = cities[city_names[end]][1], cities[city_names[end]][0]
    m.drawgreatcircle(start_lon, start_lat, end_lon, end_lat, color='blue', linewidth=2)

# Add a title
plt.title(f"Shortest Path from {city_names[start_city]} to {city_names[end_city]}")

# Display the map
plt.show()

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import time
from datetime import datetime
import random
from IPython.display import display

# Set the aesthetic style for plots
sns.set(style="whitegrid")
plt.rcParams.update({'font.size': 12})

# Function to simulate ping to different servers
def simulate_ping(server, num_pings=10):
    """Simulate ping times to a server with realistic patterns"""
    # Base latency depends on server location
    base_latency = {
        'Local': 5,
        'Regional': 25,
        'International': 120,
        'Cloud-A': 15,
        'Cloud-B': 45
    }
    
    # Simulate variation in ping times
    base = base_latency[server]
    jitter = np.random.normal(0, base * 0.1, num_pings)  # 10% jitter
    packet_loss = np.random.choice([1, 0], num_pings, p=[0.05, 0.95])  # 5% packet loss
    
    # Add occasional network congestion spikes
    congestion = np.zeros(num_pings)
    if random.random() < 0.3:  # 30% chance of congestion during the test
        congestion_start = random.randint(0, num_pings - 3)
        congestion_duration = random.randint(1, 3)
        congestion[congestion_start:congestion_start+congestion_duration] = base * 0.5
    
    # Calculate ping times
    ping_times = base + jitter + congestion
    
    # Apply packet loss (set to NaN)
    ping_times[packet_loss == 1] = np.nan
    
    return ping_times

# Simulate data collection
def collect_network_data(servers, samples_per_server=60):
    """Collect simulated network data for multiple servers"""
    data = []
    
    print("Collecting network data...")
    for server in servers:
        print(f"Pinging {server}...")
        for i in range(samples_per_server // 10):
            # Simulate 10 pings at a time
            ping_results = simulate_ping(server, 10)
            timestamp = datetime.now()
            
            for j, ping in enumerate(ping_results):
                data.append({
                    'timestamp': timestamp,
                    'server': server,
                    'ping_ms': ping,
                    'packet_number': i * 10 + j + 1
                })
            
            # Simulate data collection delay
            time.sleep(0.1)
    
    return pd.DataFrame(data)

# Analyze network data
def analyze_network_data(df):
    """Perform analysis on network data"""
    # Calculate summary statistics
    summary = df.groupby('server')['ping_ms'].agg([
        ('avg_ping', 'mean'),
        ('min_ping', 'min'),
        ('max_ping', 'max'),
        ('std_ping', 'std'),
        ('packet_loss', lambda x: x.isna().mean() * 100)  # Calculate packet loss percentage
    ]).round(2)
    
    # Calculate jitter (variation in ping times)
    jitter = df.groupby('server')['ping_ms'].diff().abs().groupby(df['server']).mean().round(2)
    summary['jitter'] = jitter
    
    # Identify servers with potential issues
    summary['status'] = 'Good'
    summary.loc[summary['packet_loss'] > 1, 'status'] = 'Warning'
    summary.loc[summary['packet_loss'] > 5, 'status'] = 'Critical'
    summary.loc[summary['avg_ping'] > 100, 'status'] = 'High Latency'
    
    return summary

# Visualize network data
def visualize_network_data(df, summary):
    """Create visualizations for network data analysis"""
    # Create a figure with subplots
    fig = plt.figure(figsize=(20, 16))
    
    # 1. Ping time series
    ax1 = plt.subplot(2, 2, 1)
    for server in df['server'].unique():
        server_data = df[df['server'] == server]
        ax1.plot(server_data['packet_number'], server_data['ping_ms'], 'o-', label=server, alpha=0.7)
    
    ax1.set_title('Ping Times Over Time')
    ax1.set_xlabel('Packet Number')
    ax1.set_ylabel('Ping (ms)')
    ax1.legend()
    ax1.grid(True)
    
    # 2. Box plot of ping times
    ax2 = plt.subplot(2, 2, 2)
    sns.boxplot(x='server', y='ping_ms', data=df, ax=ax2)
    ax2.set_title('Distribution of Ping Times')
    ax2.set_xlabel('Server')
    ax2.set_ylabel('Ping (ms)')
    
    # 3. Packet loss bar chart
    ax3 = plt.subplot(2, 2, 3)
    sns.barplot(x=summary.index, y='packet_loss', data=summary, ax=ax3)
    ax3.set_title('Packet Loss by Server')
    ax3.set_xlabel('Server')
    ax3.set_ylabel('Packet Loss (%)')
    
    # 4. Average ping with error bars
    ax4 = plt.subplot(2, 2, 4)
    sns.barplot(x=summary.index, y='avg_ping', data=summary, ax=ax4)
    
    # Add error bars representing standard deviation
    for i, server in enumerate(summary.index):
        ax4.errorbar(i, summary.loc[server, 'avg_ping'], 
                     yerr=summary.loc[server, 'std_ping'], 
                     fmt='none', color='black', capsize=5)
    
    ax4.set_title('Average Ping Time by Server')
    ax4.set_xlabel('Server')
    ax4.set_ylabel('Average Ping (ms)')
    
    plt.tight_layout()
    return fig

# Run the network analysis
def run_network_analysis():
    """Run complete network analysis workflow"""
    # Define servers to test
    servers = ['Local', 'Regional', 'International', 'Cloud-A', 'Cloud-B']
    
    # Collect data
    df = collect_network_data(servers, samples_per_server=60)
    
    # Display raw data sample
    print("\nRaw data sample:")
    display(df.head())
    
    # Analyze data
    summary = analyze_network_data(df)
    
    # Display summary
    print("\nNetwork Analysis Summary:")
    display(summary)
    
    # Create visualizations
    fig = visualize_network_data(df, summary)
    plt.show()
    
    # Additional analysis: correlation heatmap
    plt.figure(figsize=(10, 8))
    pivot_data = df.pivot_table(
        index='packet_number', 
        columns='server', 
        values='ping_ms'
    )
    
    # Calculate and plot correlation matrix
    corr = pivot_data.corr()
    sns.heatmap(corr, annot=True, cmap='coolwarm', vmin=-1, vmax=1, center=0)
    plt.title('Correlation of Ping Times Between Servers')
    plt.tight_layout()
    plt.show()
    
    return df, summary

# Run the analysis
network_data, network_summary = run_network_analysis()

import numpy as np
import matplotlib.pyplot as plt

def calculate_trajectory(v0, angle_deg, g=9.81):
    """
    Calculate the trajectory of a projectile
    
    Parameters:
    v0 (float): Initial velocity in m/s
    angle_deg (float): Launch angle in degrees
    g (float): Acceleration due to gravity in m/s²
    """
    # Convert angle to radians
    angle = np.radians(angle_deg)
    
    # Initial velocities
    v0x = v0 * np.cos(angle)
    v0y = v0 * np.sin(angle)
    
    # Time of flight
    t_flight = 2 * v0y / g
    
    # Create time array
    t = np.linspace(0, t_flight, 100)
    
    # Calculate x and y positions
    x = v0x * t
    y = v0y * t - 0.5 * g * t**2
    
    # Calculate maximum height and range
    h_max = v0y**2 / (2 * g)
    range_x = v0x * t_flight
    
    return t, x, y, h_max, range_x

# Set initial conditions
v0 = 20  # Initial velocity (m/s)
angles = [30, 45, 60]  # Launch angles (degrees)

# Create subplot
plt.figure(figsize=(12, 6))

# Plot trajectories for different angles
for angle in angles:
    t, x, y, h_max, range_x = calculate_trajectory(v0, angle)
    plt.plot(x, y, label=f'Angle = {angle}°')
    
    # Print results
    print(f"\nFor {angle}° launch angle:")
    print(f"Maximum height: {h_max:.2f} meters")
    print(f"Range: {range_x:.2f} meters")

plt.title('Projectile Motion for Different Launch Angles')
plt.xlabel('Distance (m)')
plt.ylabel('Height (m)')
plt.grid(True)
plt.legend()
plt.axis('equal')
plt.show()

import numpy as np
import matplotlib.pyplot as plt

# Define the luminosity range (in terms of the Sun's luminosity)
luminosities = np.linspace(0.1, 10, 100)  # From 0.1 to 10 times the Sun's luminosity

# Compute the habitable zone boundaries
inner_hz = np.sqrt(luminosities / 1.1)
outer_hz = np.sqrt(luminosities / 0.53)

# Plot the habitable zone
plt.figure(figsize=(8, 5))
plt.plot(luminosities, inner_hz, label='Inner Boundary', color='red')
plt.plot(luminosities, outer_hz, label='Outer Boundary', color='blue')
plt.fill_between(luminosities, inner_hz, outer_hz, color='green', alpha=0.3, label='Habitable Zone')

# Labels and title
plt.xlabel("Luminosity (L/L_sun)")
plt.ylabel("Distance (AU)")
plt.title("Habitable Zone Boundaries for Different Star Luminosities")
plt.legend()
plt.grid()
plt.show()

import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
import matplotlib.pyplot as plt

# Generate sample data
np.random.seed(42)
X = 2 * np.random.rand(100, 1)
y = 4 + 3 * X + np.random.randn(100, 1)

# Split the data into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Create and train the model
model = LinearRegression()
model.fit(X_train, y_train)

# Make predictions
y_pred = model.predict(X_test)

# Print model coefficients
print(f"Intercept: {model.intercept_[0]:.2f}")
print(f"Slope: {model.coef_[0][0]:.2f}")

# Calculate R-squared score
r2_score = model.score(X_test, y_test)
print(f"R-squared score: {r2_score:.2f}")

# Visualization
plt.figure(figsize=(10, 6))
plt.scatter(X_train, y_train, color='blue', label='Training Data')
plt.scatter(X_test, y_test, color='green', label='Testing Data')
plt.plot(X_test, y_pred, color='red', linewidth=2, label='Regression Line')
plt.xlabel('X')
plt.ylabel('y')
plt.title('Linear Regression Example')
plt.legend()
plt.grid(True)
plt.show()

import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
import seaborn as sns

# Generate sample customer transaction data
np.random.seed(42)

# Create sample data
n_customers = 1000
customer_ids = range(1, n_customers + 1)

# Generate random dates within the last year
end_date = datetime(2025, 2, 21)
start_date = end_date - timedelta(days=365)
dates = [start_date + timedelta(days=np.random.randint(0, 365)) for _ in range(n_customers)]

# Generate random transaction amounts and frequencies
frequencies = np.random.poisson(lam=3, size=n_customers)
monetary_values = np.random.normal(100, 30, n_customers)
monetary_values = [max(20, val) for val in monetary_values]  # Ensure minimum transaction value

# Create DataFrame
data = pd.DataFrame({
    'customer_id': customer_ids,
    'last_purchase_date': dates,
    'purchase_frequency': frequencies,
    'total_monetary_value': monetary_values
})

# Calculate RFM scores
current_date = end_date
data['recency'] = (current_date - pd.to_datetime(data['last_purchase_date'])).dt.days

# Function to assign RFM scores
def assign_score(value, quartiles):
    if value <= quartiles[0.25]:
        return 4
    elif value <= quartiles[0.5]:
        return 3
    elif value <= quartiles[0.75]:
        return 2
    else:
        return 1

# Calculate quartiles and assign scores
r_quartiles = data['recency'].quantile([0.25, 0.5, 0.75])
f_quartiles = data['purchase_frequency'].quantile([0.25, 0.5, 0.75])
m_quartiles = data['total_monetary_value'].quantile([0.25, 0.5, 0.75])

data['R_score'] = data['recency'].apply(lambda x: assign_score(x, r_quartiles))
data['F_score'] = data['purchase_frequency'].apply(lambda x: 5 - assign_score(x, f_quartiles))
data['M_score'] = data['total_monetary_value'].apply(lambda x: 5 - assign_score(x, m_quartiles))

# Calculate final RFM score
data['RFM_score'] = data['R_score'] * 100 + data['F_score'] * 10 + data['M_score']

# Segment customers based on RFM score
def segment_customers(row):
    if row['RFM_score'] >= 444:
        return 'Best Customers'
    elif row['RFM_score'] >= 334:
        return 'Loyal Customers'
    elif row['RFM_score'] >= 224:
        return 'Average Customers'
    else:
        return 'Lost Customers'

data['customer_segment'] = data.apply(segment_customers, axis=1)

# Create visualizations
plt.figure(figsize=(15, 10))

# Plot 1: Customer Segments Distribution
plt.subplot(2, 2, 1)
segment_counts = data['customer_segment'].value_counts()
plt.pie(segment_counts, labels=segment_counts.index, autopct='%1.1f%%')
plt.title('Customer Segments Distribution')

# Plot 2: Average Monetary Value by Segment
plt.subplot(2, 2, 2)
avg_monetary = data.groupby('customer_segment')['total_monetary_value'].mean()
avg_monetary.plot(kind='bar')
plt.title('Average Monetary Value by Segment')
plt.ylabel('Average Spend ($)')
plt.xticks(rotation=45)

# Plot 3: Recency vs Frequency
plt.subplot(2, 2, 3)
plt.scatter(data['recency'], data['purchase_frequency'], 
           c=data['total_monetary_value'], cmap='viridis')
plt.colorbar(label='Monetary Value ($)')
plt.xlabel('Recency (days)')
plt.ylabel('Purchase Frequency')
plt.title('Recency vs Frequency (color = Monetary Value)')

# Plot 4: RFM Score Distribution
plt.subplot(2, 2, 4)
plt.hist(data['RFM_score'], bins=20)
plt.xlabel('RFM Score')
plt.ylabel('Number of Customers')
plt.title('RFM Score Distribution')

plt.tight_layout()
plt.show()

# Print summary statistics
print("\nCustomer Segment Summary:")
print(data.groupby('customer_segment').agg({
    'customer_id': 'count',
    'total_monetary_value': 'mean',
    'purchase_frequency': 'mean',
    'recency': 'mean'
}).round(2))