"""Shared scientific plots of retained fields; no solver or synthetic fields."""
import argparse
import json
from pathlib import Path
import sys

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm
from matplotlib.patches import Circle,Rectangle

ROOT=Path(__file__).resolve().parents[2]
sys.path.insert(0,str(ROOT))
from research.continuous_viv.render import vorticity


def decorate(ax,scene,role,y,*,mean=False):
    if not mean:
        if scene=='karman':ax.add_patch(Circle((-20,0),1,facecolor='white',edgecolor='#253544',lw=.7))
        if role!='target':
            for x,dy in ((0,0),(1.3,.75),(1.3,-.75)):
                ax.add_patch(Circle((x,dy+y),.5,facecolor='white',edgecolor='#253544',lw=.7))
    ax.add_patch(Rectangle((4,-5),20,10,fill=False,ls='--',lw=.6,edgecolor='#555'))
    ax.set_xlim(-23.6,27.6);ax.set_ylim(-6.4,6.4)
    ax.set_xlabel('(x - pinball-front equilibrium centre) / D',fontsize=8)
    ax.set_ylabel('(y - channel centre) / D',fontsize=8)
    ax.tick_params(labelsize=8)


def main():
    ap=argparse.ArgumentParser(description=__doc__);ap.add_argument('directory',type=Path)
    args=ap.parse_args();p=args.directory
    case=json.loads((p/'case.json').read_text());scene=case['scene']
    cmap=plt.get_cmap('RdBu_r',21).copy();cmap.set_bad('#DDE2E5')
    norm=BoundaryNorm(np.linspace(-3,3,22),21,clip=True)
    prefix='SMOKE — plumbing only | ' if case['plan']['smoke'] else ''
    roles=('target','zero','controlled')
    data={};motion={}
    for role in roles:
        with np.load(p/role/'fields.npz') as z:data[role]={k:z[k] for k in z.files}
        with np.load(p/role/'native.npz') as z:
            motion[role]=z['rows'][:,1].copy()/40.
    shape=data['target']['ux'].shape[-2:]
    extent=((-0.5-1200)/40,(shape[1]-.5-1200)/40,
            (-.5-510.5)/40,(shape[0]-.5-510.5)/40)
    fig,axs=plt.subplots(3,1,figsize=(13,10),constrained_layout=True)
    for ax,role in zip(axs,roles):
        d=data[role];flags=d['common_fluid_all_boundaries'].astype(np.uint16)
        w=vorticity(d['fixed_window_mean_ux'],d['fixed_window_mean_uy'],flags)
        image=ax.imshow(w,origin='lower',extent=extent,cmap=cmap,norm=norm,interpolation='nearest')
        decorate(ax,scene,role,0,mean=True)
        label='no pinball reference' if role=='target' else ('zero rotation' if role=='zero' else ('constant rotations' if scene=='steady' else 'frozen PPO'))
        ax.set_title(f'{role}: {label} | arithmetic fixed-window mean',fontsize=10)
    fig.suptitle(prefix+f'{scene}: modern Re_D=50, m*=4, zero structural damping\n'
                 'Gray: cells outside the common fluid support; mean fields are not instantaneous geometry',fontsize=12)
    fig.colorbar(image,ax=list(axs),shrink=.65,label='vorticity D/U, fixed limits [-3,3]')
    fig.savefig(p/'mean-flow.png',dpi=160);plt.close(fig)
    slots=max(len(data[r]['ux']) for r in roles)
    fig,axs=plt.subplots(3,slots,figsize=(max(8,3*slots),7),squeeze=False,constrained_layout=True)
    for row,role in enumerate(roles):
        d=data[role]
        for col,ax in enumerate(axs[row]):
            if col>=len(d['ux']):ax.axis('off');continue
            w=vorticity(d['ux'][col],d['uy'][col],d['flags'][col])
            image=ax.imshow(w,origin='lower',extent=extent,cmap=cmap,norm=norm,interpolation='nearest')
            step=int(d['selected_response_steps'][col]);y=float(motion[role][step-1])
            decorate(ax,scene,role,y)
            ax.set_xlim(-3,14);ax.set_ylim(-3,3)
            ax.set_xlabel('x/D relative to equilibrium front',fontsize=7)
            ax.set_title(f'{role} | slot {col+1}\ntU/D={step*.005/40:.2f}',fontsize=8)
    family={'karman':'independently phased diagnostic snapshots (different real times)',
            'steady':'late snapshots at matching real times',
            'taylor':'common event-relative snapshots at matching real times'}[scene]
    fig.suptitle(prefix+f'{scene}: {family}\nFixed lab coordinates; same vorticity normalization/limits; source '+case['source'][:8],fontsize=12)
    fig.colorbar(image,ax=list(axs.ravel()),shrink=.7,label='vorticity D/U [-3,3]')
    fig.savefig(p/'selected-flow.png',dpi=160);plt.close(fig)
    fig,ax=plt.subplots(figsize=(11,3.5),constrained_layout=True)
    for role,color in (('zero','#5A6570'),('controlled','#147D73')):
        y=motion[role];stride=max(1,len(y)//12000)
        time=(np.arange(len(y))+1)*.005/40
        ax.plot(time[::stride],y[::stride],lw=.8,label=role,color=color)
    ax.set(xlabel='response time tU/D',ylabel='transverse y/D',ylim=(-1,1))
    ax.legend();ax.grid(alpha=.2)
    ax.set_title(prefix+scene+': independent paired releases; full retained duration')
    fig.savefig(p/'motion.png',dpi=160);plt.close(fig)
    print(json.dumps(dict(scene=scene,plots=['mean-flow.png','selected-flow.png','motion.png'],
                          vorticity_limits=[-3,3],normalization='curl(u) D/U; five-node fluid stencil',
                          mean_mask='intersection over collected boundary fields',display_interpolation='nearest')))


if __name__=='__main__':main()