"""课程 MECH5961_TMM_design.m 的 Python 移植；长度输入单位为 mm。"""
from functools import lru_cache
import numpy as np
from scipy.special import jv, jve, jn_zeros

AIR = dict(rho=1.204, pressure=101325., gamma=1.4, viscosity=1.81e-5,
           heat_capacity=1005., conductivity=.0257)
BASELINE = [dict(name="Type A", count=8, depth=20., radius=.70, extension=2.),
            dict(name="Type B", count=8, depth=20., radius=.95, extension=4.)]


def validate(cells):
    if not isinstance(cells, list) or not 1 <= len(cells) <= 3:
        raise ValueError("使用 1–3 种腔体。")
    names = []
    for cell in cells:
        name = str(cell['name']).strip()
        if not name or name in names:
            raise ValueError("腔体类型名称不可为空或重复。")
        names.append(name)
        for key in ('count', 'depth', 'radius', 'extension'):
            if not np.isfinite(float(cell[key])):
                raise ValueError("参数必须为有限数值。")
        if float(cell['count']) != int(cell['count']) or cell['count'] < 1:
            raise ValueError("腔体数量必须为正整数。")
        if not 10 <= cell['depth'] <= 20:
            raise ValueError("腔深必须为 10–20 mm。")
        if not .5 <= cell['radius'] <= 1.25:
            raise ValueError("孔半径必须为 0.50–1.25 mm。")
        if not 0 <= cell['extension'] <= 10 or cell['extension'] >= cell['depth']:
            raise ValueError("颈伸长须为 0–10 mm 且小于腔深。")
    if sum(c['count'] for c in cells) != 16:
        raise ValueError("腔体总数必须恰好为 16。")


@lru_cache(maxsize=2)
def roots(terms):
    return jn_zeros(1, terms)


def end_correction(radius, receiving_radius, terms):
    r = roots(terms)
    x = r * radius / receiving_radius
    return 4 * radius * np.sum(jv(1, x)**2 / (x * (r * jv(0, r))**2))


def loss(omega, diffusivity, radius):
    q = radius * np.sqrt(-1j * omega / diffusivity)
    small = np.abs(q) < 1e-3
    safe = np.where(small, 1.+0j, q)
    exact = 1 - 2 * jve(1, safe) / (safe * jve(0, safe))
    series = -q**2/8 - q**4/48 - 11*q**6/3072
    return np.where(small, series, exact)


def medium(omega, radius, area):
    a = AIR
    density = a['rho'] / loss(omega, a['viscosity']/a['rho'], radius)
    compressibility = (a['gamma']-(a['gamma']-1)*loss(
        omega, a['conductivity']/(a['rho']*a['heat_capacity']), radius))/(a['gamma']*a['pressure'])
    return np.sqrt(density/compressibility)/area, omega*np.sqrt(density*compressibility)


def cell_impedance(cell, f, terms):
    depth, radius, ext = [cell[k]*1e-3 for k in ('depth', 'radius', 'extension')]
    area = .011**2
    viscous_radius = 2*area/(4*.011)
    length = .0015 + ext + end_correction(radius, np.sqrt(.00015625/np.pi), terms)
    length += end_correction(radius, np.sqrt(area/np.pi), terms)
    zn, kn = medium(2*np.pi*f, radius, np.pi*radius**2)
    zc, kc = medium(2*np.pi*f, viscous_radius, area)
    admittance = 1j*np.tan(kc*(depth-ext))/zc
    if ext > 0:
        surrounding = area - np.pi*(radius+.0006)**2
        admittance += 1j*np.tan(kc*ext)/(zc*area/surrounding)
    co, si = np.cos(kn*length), np.sin(kn*length)
    zin = (co + 1j*zn*si*admittance)/(1j*si/zn + co*admittance)
    c = np.sqrt(AIR['gamma']*AIR['pressure']/AIR['rho'])
    z = zin*area/(AIR['rho']*c)
    if not np.all(np.isfinite(z)) or np.any(z.real < -1e-9):
        raise ValueError("出现非有限值或非被动阻抗。")
    return z


def band_average(f, alpha, lo=800., hi=1200.):
    if not f[0] <= lo < hi <= f[-1]:
        raise ValueError("计算频率必须覆盖完整统计频段。")
    interior = (f > lo) & (f < hi)
    x = np.r_[lo, f[interior], hi]
    y = np.r_[np.interp(lo, f, alpha), alpha[interior], np.interp(hi, f, alpha)]
    return float(np.trapezoid(y, x)/(hi-lo))


def calculate(cells=None, refined=False):
    cells = BASELINE if cells is None else cells
    validate(cells)
    step, terms = (.25, 80) if refined else (1., 40)
    f = np.arange(200., 2500.+step/2, step)
    y = sum(cell['count']*.011**2/.0025/cell_impedance(cell, f, terms) for cell in cells)
    r = (1-y)/(1+y)
    alpha = 1-abs(r)**2
    if not np.all(np.isfinite(alpha)) or np.any(alpha < -1e-9) or np.any(alpha > 1+1e-9):
        raise ValueError("吸声系数超出被动范围。")
    alpha = np.clip(alpha, 0, 1)
    peak = int(np.argmax(alpha))
    return dict(frequency=f.tolist(), alpha=alpha.tolist(), band_mean=band_average(f, alpha),
                peak_frequency=float(f[peak]), peak_alpha=float(alpha[peak]),
                step_hz=step, end_terms=terms, cells=cells)


if __name__ == '__main__':
    import argparse, json
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--cells', help='包含 1–3 种腔体参数的 JSON 文件')
    parser.add_argument('--refined', action='store_true')
    parser.add_argument('--output', default='absorption.csv')
    args = parser.parse_args()
    cells = json.load(open(args.cells, encoding='utf-8')) if args.cells else BASELINE
    result = calculate(cells, args.refined)
    np.savetxt(args.output, np.column_stack([result['frequency'], result['alpha']]),
               delimiter=',', header='frequency_hz,absorption', comments='')
    print(json.dumps({k:v for k,v in result.items() if k not in ('frequency','alpha')}, ensure_ascii=False, indent=2))