import numpy as np, pandas as pd
from pathlib import Path
from scipy.interpolate import griddata
from scipy.ndimage import gaussian_filter
from scipy.signal import detrend
from numpy.linalg import lstsq
OUT=Path('/mnt/data/uruana_v4_work/source_outputs'); OUT.mkdir(parents=True,exist_ok=True)
GRID='/mnt/data/uruana_v4_source/0032_SAO_FRANCISCO_MAG_grid.asc'
x0,y0=365381.6617,8224096.1380

# Read grid
d=pd.read_csv(GRID, comment='/', na_values='*', encoding='latin1')
for c in ['X','Y','Mag','mag_var']: d[c]=pd.to_numeric(d[c],errors='coerce')
xs=np.sort(d.X.unique()); ys=np.sort(d.Y.unique()); dx=np.median(np.diff(xs)); dy=np.median(np.diff(ys))
piv=d.pivot(index='Y',columns='X',values='mag_var').reindex(index=ys,columns=xs)
Z=piv.to_numpy(float)
print('grid',Z.shape,'dxdy',dx,dy,'nan',np.isnan(Z).mean())

# Helpers
def poly_detrend(X,Y,Z,degree=1):
    m=np.isfinite(Z)
    xx=X[m]; yy=Y[m]; zz=Z[m]
    terms=[np.ones_like(xx),xx,yy]
    if degree>=2: terms += [xx*xx,xx*yy,yy*yy]
    A=np.vstack(terms).T
    coef=lstsq(A,zz,rcond=None)[0]
    terms_full=[np.ones_like(X),X,Y]
    if degree>=2: terms_full += [X*X,X*Y,Y*Y]
    trend=sum(c*t for c,t in zip(coef,terms_full))
    return Z-trend

def fill_nearest(A):
    if not np.isnan(A).any(): return A.copy()
    yy,xx=np.indices(A.shape); m=np.isfinite(A)
    return griddata((xx[m],yy[m]),A[m],(xx,yy),method='nearest')

def radial_spectrum(radius_km,degree=1):
    # square window halfwidth radius
    ix=np.where(np.abs(xs-x0)<=radius_km*1000)[0]; iy=np.where(np.abs(ys-y0)<=radius_km*1000)[0]
    zz=Z[np.ix_(iy,ix)]; x=xs[ix]-x0; y=ys[iy]-y0
    X,Y=np.meshgrid(x,y)
    zz=poly_detrend(X/1000,Y/1000,zz,degree)
    zz=fill_nearest(zz)
    wy=np.hanning(len(y)); wx=np.hanning(len(x)); W=wy[:,None]*wx[None,:]
    F=np.fft.fft2(zz*W)
    P=np.abs(F)**2
    kx=2*np.pi*np.fft.fftfreq(len(x),d=dx/1000) # rad/km
    ky=2*np.pi*np.fft.fftfreq(len(y),d=dy/1000)
    KX,KY=np.meshgrid(kx,ky); K=np.sqrt(KX*KX+KY*KY)
    # radial bins one fundamental interval
    dk=max(2*np.pi/(len(x)*dx/1000),2*np.pi/(len(y)*dy/1000))
    edges=np.arange(0,K.max()+dk,dk)
    kk=[]; lp=[]; n=[]
    for a,b in zip(edges[:-1],edges[1:]):
        m=(K>=a)&(K<b)&(K>0)
        vals=P[m]
        if vals.size>=4 and np.nanmean(vals)>0:
            kk.append(np.nanmean(K[m])); lp.append(np.log(np.nanmean(vals))); n.append(vals.size)
    kk=np.array(kk); lp=np.array(lp); n=np.array(n)
    # choose two-segment fit among reasonable splits, excluding first DC-near and top noisy Nyquist tail
    valid=(kk>0.05)&(kk<1.8) & np.isfinite(lp)
    kv=kk[valid]; lv=lp[valid]
    best=None
    # sorted already
    for j in range(4,len(kv)-4):
        # low-k deep and high-k shallow; linear each
        c1=np.polyfit(kv[:j],lv[:j],1); c2=np.polyfit(kv[j:],lv[j:],1)
        pred=np.r_[np.polyval(c1,kv[:j]),np.polyval(c2,kv[j:])]
        sse=np.sum((lv-pred)**2)
        if c1[0]<0 and c2[0]<0 and (best is None or sse<best[0]): best=(sse,j,c1,c2)
    if best is None: return None
    sse,j,c1,c2=best
    zdeep=-c1[0]/2; zsh=-c2[0]/2
    # r2s
    def r2(x,y,c):
        pr=np.polyval(c,x); return 1-np.sum((y-pr)**2)/np.sum((y-y.mean())**2)
    return dict(radius_km=radius_km,detrend_degree=degree,nx=len(x),ny=len(y),split_k_rad_km=kv[j],deep_apparent_km=zdeep,shallow_apparent_km=zsh,deep_r2=r2(kv[:j],lv[:j],c1),shallow_r2=r2(kv[j:],lv[j:],c2),sse=sse)

spec=[]
for r in [20,25,30,35,40,45,50]:
  for deg in [1,2]:
    q=radial_spectrum(r,deg)
    if q: spec.append(q)
spec=pd.DataFrame(spec); spec.to_csv(OUT/'2008_radial_spectrum_sensitivity.csv',index=False)
print('\nSPECTRUM')
print(spec.to_string(index=False))
print('shallow median',spec.shallow_apparent_km.median(),'range',spec.shallow_apparent_km.min(),spec.shallow_apparent_km.max())
print('deep median',spec.deep_apparent_km.median(),'range',spec.deep_apparent_km.min(),spec.deep_apparent_km.max())

# Euler deconvolution on local 60km square, fft derivatives; use filled/detrended grid.
# standard equation: (x-x0)Fx+(y-y0)Fy+(z-z0)Fz=N(B-F), solve x0,y0,z0,B.
def upward_and_derivs(radius_km=45, degree=1, upward_km=0):
    ix=np.where(np.abs(xs-x0)<=radius_km*1000)[0]; iy=np.where(np.abs(ys-y0)<=radius_km*1000)[0]
    zz=Z[np.ix_(iy,ix)]; xv=xs[ix]/1000; yv=ys[iy]/1000
    X,Y=np.meshgrid(xv,yv)
    zz=poly_detrend(X-X.mean(),Y-Y.mean(),zz,degree); zz=fill_nearest(zz)
    ny,nx=zz.shape
    kx=2*np.pi*np.fft.fftfreq(nx,d=dx/1000); ky=2*np.pi*np.fft.fftfreq(ny,d=dy/1000)
    KX,KY=np.meshgrid(kx,ky); K=np.sqrt(KX*KX+KY*KY)
    F=np.fft.fft2(zz)
    if upward_km>0: F=F*np.exp(-K*upward_km)
    field=np.real(np.fft.ifft2(F))
    Fx=np.real(np.fft.ifft2(1j*KX*F)); Fy=np.real(np.fft.ifft2(1j*KY*F)); Fz=np.real(np.fft.ifft2(K*F)) # downward continuation derivative convention amplitude
    return X,Y,field,Fx,Fy,Fz

def euler_family(N,win,up):
    X,Y,F,Fx,Fy,Fz=upward_and_derivs(45,1,up)
    nr,nc=F.shape; h=win//2; sols=[]
    # decimate centers every 2 cells
    for i in range(h,nr-h,2):
      for j in range(h,nc-h,2):
        sl=(slice(i-h,i+h+1),slice(j-h,j+h+1))
        xx=X[sl].ravel(); yy=Y[sl].ravel(); ff=F[sl].ravel(); fxx=Fx[sl].ravel(); fyy=Fy[sl].ravel(); fzz=Fz[sl].ravel()
        # equation x0 Fx + y0 Fy + z0 Fz + N B = x Fx + y Fy + z_obs Fz + N F ; z_obs=0
        A=np.c_[fxx,fyy,fzz,np.full_like(ff,N)]
        rhs=xx*fxx+yy*fyy+N*ff
        # N=0 cannot solve B; use 3 cols
        if N==0:
            A=np.c_[fxx,fyy,fzz]
        try:
            sol,resid,rank,s=np.linalg.lstsq(A,rhs,rcond=None)
        except: continue
        if rank<A.shape[1]: continue
        xsrc,ysrc,zsrc=sol[:3]
        # our Fz sign convention yields zsrc? accept positive depth as positive zsrc after testing; take abs only no
        pred=A@sol; rmse=np.sqrt(np.mean((pred-rhs)**2))
        # source within 20 km of Uruana and plausible 0-20km depth
        dist=np.hypot(xsrc-x0/1000,ysrc-y0/1000)
        depth=zsrc
        if np.isfinite(depth) and 0<depth<20 and dist<20:
            sols.append((xsrc,ysrc,depth,dist,rmse))
    if not sols: return None
    a=np.array(sols)
    # retain lower 75% residual to avoid poor windows
    q=a[:,4]<=np.quantile(a[:,4],.75); a=a[q]
    return dict(SI=N,window_cells=win,upward_km=up,n=len(a),depth_median_km=np.median(a[:,2]),depth_q25_km=np.quantile(a[:,2],.25),depth_q75_km=np.quantile(a[:,2],.75),source_dist_median_km=np.median(a[:,3]),rmse_median=np.median(a[:,4]))

eu=[]
for up in [0,1.5,3.0]:
 for N in [0,0.5,1,2]:
  for w in [5,7,9]:
   q=euler_family(N,w,up)
   if q: eu.append(q)
eu=pd.DataFrame(eu); eu.to_csv(OUT/'2008_euler_sensitivity.csv',index=False)
print('\nEULER')
print(eu.to_string(index=False))

# 2011 per-profile spectral apparent depth on residual (MAGBRU - IGRF), fixed geographic windows north of center.
r=pd.read_csv('/mnt/data/uruana_v4_source/3087_target_lines_raw.csv')
for c in ['line','Y','MAGBRU','IGRF']: r[c]=pd.to_numeric(r[c],errors='coerce')
def profile_depth(g,ylo_km,yhi_km,wlo_km,whi_km):
    y=(g.Y.to_numpy()-y0)/1000; v=(g.MAGBRU-g.IGRF).to_numpy(float)
    m=(y>=ylo_km)&(y<=yhi_km)&np.isfinite(v)
    y=y[m];v=v[m]
    if len(y)<100:return np.nan,np.nan
    order=np.argsort(y);y=y[order];v=v[order]
    # bin to 100 m medians
    bins=np.arange(y.min(),y.max()+.1,.1); bi=np.digitize(y,bins); yy=[];vv=[]
    for b in np.unique(bi):
      q=bi==b
      if q.sum(): yy.append(np.median(y[q]));vv.append(np.median(v[q]))
    yy=np.array(yy);vv=np.array(vv); n=len(vv)
    # linear detrend, hanning
    vv=detrend(vv); vv*=np.hanning(n)
    F=np.fft.rfft(vv); P=np.abs(F)**2; freq=np.fft.rfftfreq(n,d=.1); k=2*np.pi*freq
    # wavelength km between wlo and whi -> k between 2pi/whi and 2pi/wlo
    mm=(k>=2*np.pi/whi_km)&(k<=2*np.pi/wlo_km)&(P>0)
    if mm.sum()<4:return np.nan,np.nan
    c=np.polyfit(k[mm],np.log(P[mm]),1); dep=-c[0]/2
    pr=np.polyval(c,k[mm]); rr=1-np.sum((np.log(P[mm])-pr)**2)/np.sum((np.log(P[mm])-np.log(P[mm]).mean())**2)
    return dep,rr
pr=[]
for ylo,yhi in [(5,45),(5,55),(5,65)]:
 for wl,wh in [(4,16),(5,20),(6,24),(8,30)]:
  vals=[]
  for ln,g in r.groupby('line'):
   z,rr=profile_depth(g,ylo,yhi,wl,wh); vals.append((ln,z,rr))
  zz=np.array([x[1] for x in vals if np.isfinite(x[1])]); rr=np.array([x[2] for x in vals if np.isfinite(x[1])])
  pr.append(dict(y_window=f'{ylo}-{yhi}',wavelength_band_km=f'{wl}-{wh}',n=len(zz),depth_median_km=np.median(zz),depth_q25_km=np.quantile(zz,.25),depth_q75_km=np.quantile(zz,.75),r2_median=np.median(rr)))
pr=pd.DataFrame(pr);pr.to_csv(OUT/'2011_profile_spectrum_sensitivity.csv',index=False)
print('\n2011 PROFILE SPECTRA')
print(pr.to_string(index=False))

# summary text
with open(OUT/'SOURCE_DEPTH_SENSITIVITY_AUDIT.txt','w') as f:
    f.write('URUANA SOURCE-DEPTH SENSITIVITY AUDIT\n')
    f.write('Purpose: test whether current public aeromagnetic products support a stable quantitative source-depth estimate before manuscript submission.\n\n')
    f.write('2008 radial power spectrum (Spector-Grant style ensemble diagnostic):\n')
    f.write(spec.to_string(index=False)); f.write('\n\n')
    f.write(f"Shallow apparent ensemble scale median={spec.shallow_apparent_km.median():.3f} km, range={spec.shallow_apparent_km.min():.3f}-{spec.shallow_apparent_km.max():.3f} km. Deep component median={spec.deep_apparent_km.median():.3f} km, range={spec.deep_apparent_km.min():.3f}-{spec.deep_apparent_km.max():.3f} km.\n")
    f.write('These are ensemble spectral scales under strong assumptions, not a unique body depth.\n\n')
    f.write('2008 Euler sensitivity:\n'); f.write(eu.to_string(index=False)); f.write('\n\n')
    f.write('Euler depth is strongly structural-index and upward-continuation dependent; no unique depth is supported.\n\n')
    f.write('2011 profile spectral sensitivity:\n'); f.write(pr.to_string(index=False));f.write('\n\n')
    f.write('Profile spectral estimates vary materially with spatial and wavelength window; no unique source depth is supported.\n\n')
    f.write('VERDICT: FAIL for a single robust source-depth estimate from current data. Do not add a precise Euler/source depth to the main paper. A ~2 km shallow spectral ensemble scale appears reproducible in the 2008 grid, but it should be described only as an apparent ensemble scale and not assigned to the Sao Domingos source without stronger source-specific convergence.\n')