import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

def combine_structure_factors(F_A, phi_A,
                              F_B, phi_B,
                              a, b,
                              plot=True):
    """
    Combine two resultant structure factors:
        F_comb = a*F1 + b*F2

    Returns:
        magnitude, phase_deg, complex_value
    """

    # Convert to complex SFs
    SF_A = F_A * np.exp(1j * np.deg2rad(phi_A))
    SF_B = F_B * np.exp(1j * np.deg2rad(phi_B))

    # Linear combination
    SF_AB_comb = a * SF_A + b * SF_B
    F_AB_comb = np.abs(SF_AB_comb)
    SF_AB_phi_A = F_AB_comb * np.exp(1j * np.deg2rad(phi_A))
    
    SF_B_extr = ( SF_AB_comb - a * SF_A ) / b
    
    
    F_AB = a * F_A + b * F_B    
    SF_AB_x8 = F_AB * np.exp(1j * np.deg2rad(phi_A))

    # x8 combination
    SF_extr_comb = (SF_AB_comb - SF_A)/b + SF_A
    
    F_extr = (F_AB_comb - F_A)/b + F_A
    SF_extr_x8 = F_extr * np.exp(1j * np.deg2rad(phi_A))

    # Optional Argand plot
    origin = 0 + 0j
    if plot:
        plt.figure(figsize=(6, 6))
        plt.arrow(origin.real, origin.imag, SF_A.real, SF_A.imag,
                  color="blue", alpha=0.6,
                  head_width=0.2, head_length=0.2,
                  length_includes_head=True, label="SFA")
        plt.arrow(origin.real, origin.imag, SF_B.real, SF_B.imag,
                  color="orange", alpha=0.6,
                  head_width=0.4, head_length=0.4,
                  length_includes_head=True, label="SFB")
        plt.arrow(origin.real, origin.imag, SF_AB_comb.real, SF_AB_comb.imag,
                  color="green", linewidth=3,
                  head_width=0.2, head_length=0.2,
                  length_includes_head=True, label="SF_extr")
        plt.arrow(origin.real, origin.imag, SF_B_extr.real, SF_B_extr.imag,
                  color="yellow", linewidth=3,
                  head_width=0.2, head_length=0.2,
                  length_includes_head=True, label="SF_B_extr")
        plt.arrow(origin.real, origin.imag, F_extr.real, F_extr.imag,
                  color="pink", linewidth=3,
                  head_width=0.2, head_length=0.2,
                  length_includes_head=True, label="F_extr")
        plt.axhline(0)
        plt.axvline(0)
        plt.axis("equal")
        plt.xlabel("Re(F)")
        plt.ylabel("Im(F)")
        plt.title(f"F_comb = aF1 + bF2  (a={a}, b={b})")
        plt.legend()
        plt.grid(True)
        plt.show()
        
    return SF_AB_comb, SF_AB_x8, SF_B_extr, SF_extr_x8


# -----------------------------
# Example usage
# -----------------------------
F_A = 100
#phi_A = 90

F_B = 50
phi_B = 0

phi_lst = np.arange(0,360,10)
b_lst = [ 0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 1 ]
df = pd.DataFrame()

for phi_A in phi_lst:
    
    for b in b_lst:
        
        a = 1 - b
        SF_AB_comb, SF_AB_x8, SF_extr_comb, SF_extr_x8 = combine_structure_factors(
            F_A, phi_A,
            F_B, phi_B,
            a, b,
            plot=True
        )
        
        data = [ {  "F_A" : F_A,
                    "phi" : phi_A,
                    "F_B" : F_B,
                    "phi_B" : phi_B,
                    "b" : b,
                    "SF_AB_comb_abs": np.abs(SF_AB_comb),
                    "SF_AB_comb_phi": np.angle(SF_AB_comb, deg=True),
                    "SF_AB_x8_abs": np.abs(SF_AB_x8),
                    "SF_AB_x8_phi": np.angle(SF_AB_x8, deg=True),
                    "SF_B_extr_abs": np.abs(SF_extr_comb),
                    "SF_B_extr_comb_phi": np.angle(SF_extr_comb, deg=True),
                    "SF_extr_x8_abs": np.abs(SF_extr_x8),
                    "SF_extr_x8_phi": np.angle(SF_extr_x8, deg=True),
                } ]
        df_1 = pd.DataFrame( data )
        df = pd.concat( ( df, df_1 ) )
        
        df = df.reset_index( drop=True )
    
print( df )

df.to_csv( "../results/complete_varied_b_and_phi.csv", sep="," )
df_mean = pd.DataFrame()
df_mean[ "b" ] = df.b   
df_mean[ "SF_extr_x8_abs" ] = ( df.groupby('b')["SF_extr_x8_abs"].transform('mean') / F_B  ) *100

print(df_mean.drop_duplicates())