import numpy as np, scipy.io.wavfile as wf, sys
from scipy.linalg import solve_toeplitz
from scipy.signal import butter, filtfilt, stft, istft

def regions(m):
    d=np.diff(m.astype(np.int8)); st=np.flatnonzero(d==1)+1; en=np.flatnonzero(d==-1)+1
    if m[0]: st=np.r_[0,st]
    if m[-1]: en=np.r_[en,len(m)]
    return st,en

def env(x,sr,ms):
    w=max(2,int(ms/1000*sr)); return np.sqrt(np.convolve(x*x,np.ones(w)/w,'same'))

def ar_interp(x,s,e,p=48):
    L=e-s; ctx=max(4*p,3*L)
    a0=max(0,s-ctx); b1=min(len(x),e+ctx)
    if s-a0<p+1 or b1-e<p+1: return False
    seg=x[a0:b1].astype(np.float64); ms=s-a0; me=e-a0; N=len(seg)
    known=np.r_[seg[:ms],seg[me:]]
    if len(known)<3*p or not np.any(known): return False
    r=np.correlate(known,known,'full')[len(known)-1:len(known)+p].astype(float)
    if r[0]<=0: return False
    r[0]*=1.0001
    try: a=solve_toeplitz((r[:p],r[:p]),r[1:p+1])
    except Exception: return False
    coef=np.r_[1.0,-a]; segz=seg.copy(); segz[ms:me]=0.0
    c=np.convolve(segz,coef,'valid'); lo=ms-p
    if lo<0 or ms+L>len(c): return False
    rhs=-np.correlate(c[lo:ms+L],coef,'valid')[:L]
    rc=np.correlate(coef,coef,'full'); col=np.zeros(L); m=min(L,p+1)
    col[:m]=rc[p:p+m]; col[0]*=1.0001
    try: sol=solve_toeplitz((col,col),rhs)
    except Exception: return False
    if not np.all(np.isfinite(sol)) or np.max(np.abs(sol))>4*np.max(np.abs(known)): return False
    x[s:e]=sol; return True

inp,outp=sys.argv[1],sys.argv[2]
sr,x=wf.read(inp); x=x.astype(np.float64)
if x.ndim>1: x=x.mean(1)
n=len(x)
DEAD=(int(15.60*sr),int(19.70*sr))

# ---------- 1. DE-CRACKLE: impulsos que vazam acima de 16kHz ----------
b,a=butter(6,16000/(sr/2),'high'); hi=filtfilt(b,a,x)
e1=env(hi,sr,1.0)                      # janela curta = localizacao precisa
db=20*np.log10(np.maximum(e1,1e-12))
st,en=regions(db>-70)
nc=0
for s,e in zip(st,en):
    if DEAD[0]<s<DEAD[1]: continue
    s2=max(0,s-int(.0008*sr)); e2=min(n,e+int(.0008*sr))
    if e2-s2<=int(.020*sr):
        if ar_interp(x,s2,e2): nc+=1
print(f"1. de-crackle      : {nc}/{len(st)} eventos de RF reparados")

# ---------- 2. DROPOUTS ----------
ed=env(x,sr,5.0); eddb=20*np.log10(np.maximum(ed,1e-12))
st,en=regions(eddb<-50); k=(en-st)>=int(.003*sr); st,en=st[k],en[k]
ok=np.ones(n,bool)
for s,e in zip(st,en): ok[max(0,s-480):min(n,e+480)]=False
blk=int(.2*sr); tone=None; bl=9e9
for i in range(0,n-blk,blk//2):
    if ok[i:i+blk].all():
        v=np.sqrt(np.mean(x[i:i+blk]**2))
        if 1e-6<v<bl: bl,tone=v,x[i:i+blk].copy()
if tone is None: tone=np.random.randn(blk)*1e-4
ni=nf=0; rec=lost=0.0
for s,e in zip(st,en):
    L=(e-s)/sr
    if L<=0.040:
        if ar_interp(x,s,e,64): ni+=1; rec+=L
    else:
        nf+=1; lost+=L; need=e-s
        fill=np.tile(tone,int(np.ceil(need/len(tone))))[:need]*0.9
        f=int(min(.008*sr,need//4))
        if f>2: fill[:f]*=np.linspace(0,1,f); fill[-f:]*=np.linspace(1,0,f)
        x[s:e]=fill
print(f"2. dropouts        : {ni} recompostos ({rec:.2f}s) | {nf} sem conteudo ({lost:.2f}s)")

# ---------- 3. REDUCAO ESPECTRAL DO CHIADO (Wiener) ----------
f,t,Z=stft(x,sr,nperseg=2048,noverlap=1536)
P=np.abs(Z)**2
tv=np.array([not(DEAD[0]<=int(tt*sr)<DEAD[1]) for tt in t])
fr=P.sum(0).copy(); fr[~tv]=np.inf
noise=np.median(P[:,np.argsort(fr)[:max(10,int(tv.sum()//40))]],1)[:,None]
OS=2.0; FLOOR=10**(-16/10)
G=np.maximum((P-OS*noise)/np.maximum(P,1e-20), FLOOR)
G=np.maximum(G, FLOOR)
# suaviza o ganho em tempo e frequencia -> evita "ruido musical"
from scipy.ndimage import uniform_filter
G=uniform_filter(G,size=(3,3))
_,y=istft(Z*G,sr,nperseg=2048,noverlap=1536)
y=y[:n]
red=10*np.log10(np.mean(x**2)/max(np.mean(y**2),1e-20))
print(f"3. reducao de ruido: chiado 4-16kHz atenuado (delta RMS {red:+.1f} dB)")

# ---------- 4. banda util: nada real acima de 15k (mp3 ja cortou em 16k) ----------
b,a=butter(8,14500/(sr/2),'low');  y=filtfilt(b,a,y)
b,a=butter(2,75/(sr/2),'high');    y=filtfilt(b,a,y)
print("4. banda           : 75 Hz - 14,5 kHz")

y=np.clip(y,-1,1); y*= 0.9/max(np.max(np.abs(y)),1e-9)
wf.write(outp,sr,y.astype(np.float32))