import numpy as np
import math

def calculate_coseismic_displacement(strike_deg, dip_deg, rake_deg, mw, 
                                     source_depth, rec_x, rec_y, mu=3e10, nu=0.25):
    """
    Calculates 3D surface coseismic displacement from an elastic point source.
    
    Parameters:
    -----------
    strike_deg, dip_deg, rake_deg : floats
        Fault orientation angles in degrees (Aki & Richards convention).
    mw : float
        Moment magnitude.
    source_depth : float
        Depth of the point source (in meters).
    rec_x, rec_y : arrays or floats
        Receiver coordinates on the surface relative to the epicenter (in meters).
    mu : float
        Shear modulus (Poisson solid default: 30 GPa).
    nu : float
        Poisson's ratio (default: 0.25).
        
    Returns:
    --------
    ux, uy, uz : arrays or floats
        Displacements in North, East, and Down directions (meters).
    """
    # 1. Calculate Seismic Moment (M0) in N-m
    m0 = 10 ** (1.5 * mw + 9.1)
    
    # 2. Convert angles to radians
    strike = np.deg2rad(strike_deg)
    dip = np.deg2rad(dip_deg)
    rake = np.deg2rad(rake_deg)

       # 3. Compute Moment Tensor components (Aki & Richards convention)
    # North (x), East (y), Down (z)
    s_st, c_st = np.sin(strike), np.cos(strike)
    s_di, c_di = np.sin(dip), np.cos(dip)
    s_ra, c_ra = np.sin(rake), np.cos(rake)

    mxx = -m0 * (s_st*c_ra*s_di + c_st*s_ra*s_di*c_di)
    myy =  m0 * (c_st*c_ra*s_di - s_st*s_ra*s_di*c_di)
    mzz =  m0 * (s_ra * s_di * c_di)
    mxy =  m0 * (c_st*c_ra*s_di - s_st*s_ra*s_di*c_di) # note: symmetric context
    mxz = -m0 * (c_st*c_ra*c_di + s_st*s_ra*(c_di**2 - s_di**2))
    myz = -m0 * (s_st*c_ra*c_di - c_st*s_ra*(c_di**2 - s_di**2))

    # 4. Geometry relative to source depth (z = 0 at surface, source at z = d)
    d = source_depth
    r2 = rec_x**2 + rec_y**2 + d**2
    r = np.sqrt(r2)
    r3 = r**3
    r5 = r**5

    # 5. Point source Green's function derivatives (Semi-infinite halfspace approximation)
    # Using standard elastic point-dislocation variables (Mindlin/Aki derivative limits at z=0)
    C = 1.0 / (4.0 * np.pi * mu)

    # Pre-calculate common factors for efficiency
    fac1 = 3.0 / r5
    fac2 = 1.0 / r3

    # We evaluate displacements (Ux, Uy, Uz) as a function of the moment components M_jk
    # multiplied by spatial Green tensor combinations at the free surface.
    # Below is the vector field projection for standard point-source tensors:

    # Terms driven by Mxx, Myy, Mzz
    ux = C * rec_x * (fac1 * (mxx*rec_x**2 + myy*rec_y**2 + mzz*d**2) - fac2*(mxx + myy + mzz))
    uy = C * rec_y * (fac1 * (mxx*rec_x**2 + myy*rec_y**2 + mzz*d**2) - fac2*(mxx + myy + mzz))
    uz = C * d * (fac1 * (mxx*rec_x**2 + myy*rec_y**2 + mzz*d**2) - fac2*(mxx + myy + mzz))

    # Off-diagonal corrections (Mxy, Mxz, Myz interactions)
    ux += C * (fac1 * 2 * mxy * rec_x**2 * rec_y + fac1 * 2 * mxz * rec_x**2 * d - fac2 * (2 * mxy * rec_y + 2 * mxz * d))
    uy += C * (fac1 * 2 * mxy * rec_x * rec_y**2 + fac1 * 2 * myz * rec_y**2 * d - fac2 * (2 * mxy * rec_x + 2 * myz * d))
    uz += C * (fac1 * 2 * mxz * rec_x * d**2 + fac1 * 2 * myz * rec_y * d**2 - fac2 * (2 * mxz * rec_x + 2 * myz * rec_y))

    return ux, uy, uz


def latlon_to_xy(lat, lon, ref_lat, ref_lon):
    """
    Converts latitude and longitude to local X, Y coordinates in meters
    relative to a given reference point (which becomes 0, 0).
    
    Parameters:
    lat (float): Latitude of the target point in decimal degrees.
    lon (float): Longitude of the target point in decimal degrees.
    ref_lat (float): Latitude of the reference point in decimal degrees.
    ref_lon (float): Longitude of the reference point in decimal degrees.
    
    Returns:
    tuple: (x, y) coordinates in meters.
    """
    # Earth radius in meters
    EARTH_RADIUS = 6371000.0  

# Convert all degrees to radians
    lat_rad = math.radians(lat)
    lon_rad = math.radians(lon)
    ref_lat_rad = math.radians(ref_lat)
    ref_lon_rad = math.radians(ref_lon)

    # Calculate differences in radians
    delta_lat = lat_rad - ref_lat_rad
    delta_lon = lon_rad - ref_lon_rad

    # Calculate X and Y coordinates
    # X relies on the cosine of the reference latitude to adjust for narrowing longitudes
    x = EARTH_RADIUS * delta_lon * math.cos(ref_lat_rad)
    y = EARTH_RADIUS * delta_lat

    return x, y


if __name__ == "__main__":
    strike, dip, rake = 279.5, 83.6, -179.3
    mw = 7.0
    source_depth = 10000.0  # meters
    print(f"{strike:6.1f} {dip:6.1f} {rake:6.1f} strike dip rake")
    
    # Define a reference point (e.g., Downtown Los Angeles)
    REFERENCE_LAT = 40.374
    REFERENCE_LON = -125.0217

    XLAT = 40.5437
    XLON = -123.8802

    x_coords, y_coords = latlon_to_xy(XLAT, XLON, REFERENCE_LAT, REFERENCE_LON)
    x = x_coords / 1000.
    y = y_coords / 1000.
    print(f"{x:6.1f} {y:6.1f} x y")


    ux, uy, uz = calculate_coseismic_displacement(strike, dip, rake, mw, 
                                     source_depth, x_coords, y_coords, mu=3e10, nu=0.25)
    ee = uy * 1000.
    nn = ux * 1000.
    uu = - uz * 1000.

    print(f"{ee:6.1f} {nn:6.1f} {uu:6.1f} e n u")