import numpy as np

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

# --- Example Usage ---
if __name__ == "__main__":
    # Define an Mw 6.0 pure strike-slip earthquake at 10km depth
    strike, dip, rake = 90.0, 90.0, 0.0
    mw = 6.0
    depth = 10000.0  # meters
    
    # Set up a grid of surface observers (from -20km to 20km)
    x = np.linspace(-20000, 20000, 5)
    y = np.linspace(-20000, 20000, 5)
    X, Y = np.meshgrid(x, y)
    
    Ux, Uy, Uz = calculate_coseismic_displacement(strike, dip, rake, mw, depth, X, Y)
    
    print("North Displacement matrix sample (meters):\n", np.round(Ux, 4))