"""双麦克风谱估计、交换校正与阻抗反演。所有通道映射显式声明。"""
import io
import numpy as np
from scipy.signal import csd, welch, get_window
from tmm import band_average
FS, NFFT, DURATION, C = 20000, 4096, 10, 343.
MISMATCH = 1.08*np.exp(1j*np.deg2rad(4))


def spectra(positions, fs=FS, nfft=NFFT):
    x = np.asarray(positions, dtype=float)
    if x.ndim != 2 or x.shape[1] != 2 or x.shape[0] < nfft or not np.isfinite(x).all():
        raise ValueError("记录必须是有限的双通道数组，且长度至少为 FFT 点数。")
    cfg = dict(fs=fs, window=get_window('hann', nfft, fftbins=True), nperseg=nfft,
               noverlap=nfft//2, nfft=nfft, detrend='constant', scaling='density')
    f, s11 = welch(x[:,0], **cfg)
    _, s22 = welch(x[:,1], **cfg)
    _, s21 = csd(x[:,0], x[:,1], **cfg)  # conj(P1) * P2
    mask = (f >= 50) & (f <= 1600)
    if np.any(s11[mask] <= 0) or np.any(s22[mask] <= 0):
        raise ValueError("分析频段内存在无信号通道。")
    h = np.divide(s21, s11, out=np.full_like(s21, np.nan), where=s11>0)
    coh = np.divide(abs(s21)**2, s11*s22, out=np.zeros_like(s11), where=s11*s22>0)
    return f, h, np.clip(coh,0,1), s11, s22


def invert(f, h, spacing=.05, distance=.1):
    k = 2*np.pi*f/C
    r = (h-np.exp(-1j*k*spacing))/(np.exp(1j*k*spacing)-h)*np.exp(2j*k*(spacing+distance))
    with np.errstate(divide='ignore', invalid='ignore'):
        z = (1+r)/(1-r)
    return r, 1-abs(r)**2, z


def process_pair(standard, switched, fs=FS, nfft=NFFT, calibration=None, fault='normal'):
    """输入均为 [通道 A, 通道 B]；标准 A 在位置1，交换后 B 在位置1。"""
    f, ha, ca, sa1, sa2 = spectra(standard, fs, nfft)
    _, hb, cb, sb1, sb2 = spectra(np.asarray(switched)[:,::-1], fs, nfft)
    # 取包含边界的邻点，供 50/1600 Hz 插值；只在有效频带附近开方和解缠绕。
    band = (f >= 40) & (f <= 1610)
    fb, a, b = f[band], ha[band], hb[band]
    ratio = a/b
    gain = np.sqrt(abs(ratio))*np.exp(.5j*np.unwrap(np.angle(ratio)))
    # 分支以校准近零幅相失配假设选择；不能只用逐点主值相位。
    if np.real(np.median(gain)) < 0:
        gain = -gain
    corrected = (a/gain+b*gain)/2
    target_f = np.r_[50., fb[(fb>50)&(fb<1600)], 1600.]
    interp = lambda z: np.interp(target_f,fb,np.real(z))+1j*np.interp(target_f,fb,np.imag(z))
    h = interp(corrected)
    coherence = np.interp(target_f,fb,(ca[band]+cb[band])/2)
    deviation = None
    if calibration is not None:
        q1, q2 = a/calibration, b*calibration
        metric = abs(q1-q2)/np.maximum(abs((q1+q2)/2),1e-12)
        deviation = np.interp(target_f,fb,metric)
    r, alpha, z = invert(target_f,h)
    max_dev = None if deviation is None else float(np.max(deviation))
    mean_coh = band_average(target_f,coherence,50,1600)
    rejected = fault=='overload' or (max_dev is not None and max_dev>.15)
    if fault=='overload': reason='本轮存在削顶过载，应降低激励后重测。'
    elif rejected: reason='独立校准校正后的成对记录不一致，教学互易性指标超过 15%。'
    elif mean_coh<.9: reason='相干性低于 0.90，需诊断并结合重复性判断，不自动拒收。'
    else: reason='本轮质量指标通过。'
    rows=[]
    for i,freq in enumerate(target_f):
        comp=lambda v: dict(re=float(v.real),im=float(v.imag))
        rows.append(dict(f=float(freq),H=comp(h[i]),R=comp(r[i]),alpha=float(alpha[i]),
                         z=comp(z[i]) if np.isfinite(z[i]) else None,
                         coherence=float(coherence[i]),deviation=None if deviation is None else float(deviation[i])))
    return dict(rows=rows,reciprocity=max_dev,coherence=mean_coh,rejected=bool(rejected),reason=reason,
                provenance=dict(engine='Python NumPy/SciPy',fs=fs,nfft=nfft,window='periodic Hann',
                overlap=nfft//2,standard_samples=len(standard),switched_samples=len(switched),
                estimator='H1 = mean(conj(P1)*P2) / mean(abs(P1)^2)',
                correction='continuous complex square root; physical positions reordered',
                reciprocity='independent calibration residual' if calibration is not None else 'not available'))


def synthetic_record(trial=0, fault='normal', swapped=False, kind='resonator'):
    """生成 200000 点实际时序；用于虚拟实验，不代表 BU/BG 性能。"""
    rng=np.random.default_rng(7100+int(trial)*17+int(swapped))
    n=FS*DURATION
    f=np.fft.rfftfreq(n,1/FS)
    # 激励覆盖正式分析频段并留少量过渡带，减少有限窗边界偏差。
    source=np.fft.rfft(rng.normal(size=n))
    envelope=np.minimum(np.clip((f-20)/30,0,1),np.clip((1800-f)/200,0,1))
    source*=envelope
    f0=960+[-9,5,12,-4,8][int(trial)%5]
    z=.85+1j*3.2*(np.maximum(f,.1)/f0-f0/np.maximum(f,.1))
    reflection=(z-1)/(z+1)
    if kind=='rigid': reflection=np.ones_like(z)
    elif kind=='matched': reflection=np.zeros_like(z)
    pressure=lambda x: np.exp(-2j*np.pi*f*x/C)+reflection*np.exp(2j*np.pi*f*x/C)
    spectra_pos=[source*pressure(-.15),source*pressure(-.1)]
    # 低相干性情境增加两路独立噪声，不把质量值硬写为常数。
    noise_level=.25 if fault=='coherence' else .0004
    for i in range(2):
        noise=np.fft.rfft(rng.normal(size=n))*envelope
        spectra_pos[i]+=noise_level*noise
    channels=spectra_pos[::-1] if swapped else spectra_pos
    channels[1]=channels[1]*MISMATCH
    if swapped and fault in ('drift','leak'):
        channels[0]=channels[0]*1.25*np.exp(.12j)
    raw=np.column_stack([np.fft.irfft(s,n=n) for s in channels])*.1
    if fault=='overload':raw=np.clip(raw,-.015,.015)
    return raw


def simulate(trial, fault, kind):
    if fault not in ('normal','coherence','drift','leak','overload') or kind not in ('resonator','rigid','matched'):
        raise ValueError('未知训练情境。')
    a=synthetic_record(trial,fault,False,kind)
    b=synthetic_record(trial,fault,True,kind)
    result=process_pair(a,b,calibration=MISMATCH,fault=fault)
    result['waveform']={'time':(np.arange(600)/FS).tolist(),'standard':a[:600].tolist(),'switched':b[:600].tolist()}
    return result


def parse_lvm(text):
    lines=text.splitlines()
    no_x=any(line.strip()=='X_Columns\tNo' for line in lines[:30])
    if no_x:
        delta=next((line.split()[1:] for line in lines[:30] if line.startswith('Delta_X')),None)
        if not delta or not all(np.isclose(float(v),1/FS,rtol=.001) for v in delta):
            raise ValueError('LVM Delta_X 不符合 20000 Hz。')
        heading=next((i for i,line in enumerate(lines) if line.startswith('X_Value')),None)
        if heading is None:raise ValueError('未找到 LVM 数据标题。')
        # X_Columns=No 的课程文件保留一个空的首列；实际数值为通道 A/B。
        data=np.loadtxt(io.StringIO('\n'.join(lines[heading+1:])),usecols=(0,1))
        expected=next((line.split()[1:] for line in lines[:30] if line.startswith('Samples')),None)
        if expected and any(int(v)!=len(data) for v in expected):raise ValueError('LVM 声明的样本数与数据不一致。')
    else:
        start=None
        for i,line in enumerate(lines):
            cols=line.strip().split()
            if len(cols)<3:continue
            try:[float(v) for v in cols[:3]]
            except ValueError:continue
            start=i;break
        if start is None:raise ValueError('未找到时间、通道 A/B 三列数值。')
        full=np.loadtxt(io.StringIO('\n'.join(lines[start:])),usecols=(0,1,2))
        if full.ndim!=2 or not np.allclose(np.diff(full[:,0]),1/FS,rtol=.02,atol=1e-8):
            raise ValueError('时间列不符合 20000 Hz。')
        data=full[:,1:3]
    if data.ndim!=2 or data.shape[1]!=2 or len(data)<NFFT or not np.isfinite(data).all():
        raise ValueError('记录长度不足或包含无效值。')
    return data


def aggregate(trials):
    if not trials:return None
    if len(trials)>30:raise ValueError('最多汇总 30 次接受试验。')
    f=np.array([r['f'] for r in trials[0]['rows']])
    if len(f)<2 or np.any(np.diff(f)<=0):raise ValueError('频率网格无效。')
    hs=[]; alphas=[]
    for trial in trials:
        if len(trial['rows'])!=len(f) or not np.allclose(f,[r['f'] for r in trial['rows']]):
            raise ValueError('各试验频率网格必须一致。')
        h=np.array([complex(r['H']['re'],r['H']['im']) for r in trial['rows']])
        if not np.isfinite(h).all():raise ValueError('传递函数包含非有限值。')
        hs.append(h);alphas.append(invert(f,h)[1])
    h=np.mean(hs,axis=0);r,a,z=invert(f,h)
    comp=lambda x:dict(re=float(x.real),im=float(x.imag))
    rows=[dict(f=float(f[i]),H=comp(h[i]),R=comp(r[i]),alpha=float(a[i]),z=comp(z[i]) if np.isfinite(z[i]) else None) for i in range(len(f))]
    bands=[]
    for lo,hi in [(400,800),(800,1200),(1200,1600)]:
        vals=[band_average(f,alpha,lo,hi) for alpha in alphas]
        mean=float(np.mean(vals));sd=float(np.std(vals,ddof=1)) if len(vals)>1 else None
        bands.append(dict(lo=lo,hi=hi,fromMeanH=band_average(f,a,lo,hi),trialMean=mean,sd=sd,cv=sd/abs(mean) if sd is not None and abs(mean)>1e-8 else None))
    return dict(rows=rows,bands=bands)