#!/usr/bin/env python3
"""Independent polar-area quadrature and eccentric-anomaly orbital reference.

Numerical constants are IAU2015 nominal GM/radius conversion factors. All
computed records are synthetic, not observed HARPS surface maps. MIT, AC18.
"""
import argparse,json
from pathlib import Path
import numpy as np
from scipy.optimize import brentq

SOLAR_GM=1.3271244e20
EARTH_GM=3.986004e14
SOLAR_RADIUS=6.957e8
AU=149597870700.
PARSEC=AU*648000/np.pi
MICROAS=648000/np.pi*1e6


def analytic_orbit(settings,time_days):
    companion=EARTH_GM*settings['planetEarthMasses'];mu=SOLAR_GM+companion
    period=settings['planetPeriodDays']*86400
    axis=(mu*(period/(2*np.pi))**2)**(1/3)
    eccentricity=settings['eccentricity']
    mean_anomaly=((2*np.pi*time_days/settings['planetPeriodDays']+np.pi)%(2*np.pi))-np.pi
    eccentric_anomaly=brentq(lambda anomaly:anomaly-eccentricity*np.sin(anomaly)-mean_anomaly,-np.pi,np.pi,xtol=1e-14)
    position=np.array([axis*(np.cos(eccentric_anomaly)-eccentricity),axis*np.sqrt(1-eccentricity**2)*np.sin(eccentric_anomaly)])
    angular_rate=2*np.pi/period/(1-eccentricity*np.cos(eccentric_anomaly))
    velocity=np.array([-axis*np.sin(eccentric_anomaly)*angular_rate,axis*np.sqrt(1-eccentricity**2)*np.cos(eccentric_anomaly)*angular_rate])
    argument=np.deg2rad(settings['periapseDegrees']);inclination=np.deg2rad(settings['orbitInclinationDegrees'])
    rotated_position=np.array([position[0]*np.cos(argument)-position[1]*np.sin(argument),position[0]*np.sin(argument)+position[1]*np.cos(argument)])
    rotated_velocity=velocity[0]*np.sin(argument)+velocity[1]*np.cos(argument)
    fraction=-companion/mu
    return dict(x_metres=fraction*rotated_position[0],y_metres=fraction*rotated_position[1]*np.cos(inclination),
                rv_m_s=fraction*rotated_velocity*np.sin(inclination))


def polar_disk(radial_order,angular_order):
    abscissa,weights=np.polynomial.legendre.leggauss(radial_order)
    radius=np.sqrt((abscissa+1)/2)[:,None]
    angle=(np.arange(angular_order)+.5)*2*np.pi/angular_order
    horizontal=np.broadcast_to(radius*np.cos(angle),(radial_order,angular_order)).ravel()
    vertical=np.broadcast_to(radius*np.sin(angle),(radial_order,angular_order)).ravel()
    toward=-np.sqrt(1-horizontal**2-vertical**2)
    area=np.broadcast_to((weights/4*2*np.pi/angular_order)[:,None],(radial_order,angular_order)).ravel()
    limb=1-.45*(1+toward)-.2*(1+toward)**2
    return horizontal,vertical,toward,area*limb


def observe(settings,time_days,radial_order=160,angular_order=640):
    horizontal,vertical,toward,quiet=polar_disk(radial_order,angular_order)
    inclination=np.deg2rad(settings['spinInclinationDegrees']);latitude=np.deg2rad(settings['spotLatitudeDegrees'])
    longitude=np.deg2rad(settings['spotLongitudeDegrees'])+2*np.pi*time_days/settings['rotationDays']
    direction=np.array([np.cos(latitude)*np.sin(longitude),
        np.sin(latitude)*np.sin(inclination)-np.cos(latitude)*np.cos(longitude)*np.cos(inclination),
        -np.sin(latitude)*np.cos(inclination)-np.cos(latitude)*np.cos(longitude)*np.sin(inclination)])
    region=np.exp(-(1-(horizontal*direction[0]+vertical*direction[1]+toward*direction[2]))/np.deg2rad(settings['spotWidthDegrees'])**2)
    if settings['spotContrast']==0 and settings['convectiveSuppressionMps']==0:region[:]=0
    weight=quiet*(1-settings['spotContrast']*region)
    orbit=analytic_orbit(settings,time_days)
    projected=2*np.pi*SOLAR_RADIUS/(settings['rotationDays']*86400)*np.sin(inclination)
    velocity=projected*horizontal+settings['convectiveSuppressionMps']*region+orbit['rv_m_s']
    total=weight.sum()
    return dict(time_days=time_days,flux=float(total/quiet.sum()),centroid_m_s=float(np.sum(weight*velocity)/total),
        reflex_velocity_m_s=float(orbit['rv_m_s']),projected_active_proxy=float(np.sum(quiet*region)/quiet.sum()),
        photocenter_x_microas=float((orbit['x_metres']+SOLAR_RADIUS*np.sum(weight*horizontal)/total)/(settings['distanceParsecs']*PARSEC)*MICROAS),
        photocenter_y_microas=float((orbit['y_metres']+SOLAR_RADIUS*np.sum(weight*vertical)/total)/(settings['distanceParsecs']*PARSEC)*MICROAS))


def main():
    parser=argparse.ArgumentParser();parser.add_argument('--input',type=Path,default=Path('surface-input.json'))
    parser.add_argument('--out',type=Path,default=Path('surface-reference.json'));args=parser.parse_args()
    inputs=json.loads(args.input.read_text());rows=[]
    for entry in inputs['cases']:
        settings=entry['settings']
        for fraction in [.03,.9/(2*np.pi),.31,.52,.73]:
            time_days=fraction*settings['rotationDays']
            coarse=observe(settings,time_days,80,320);fine=observe(settings,time_days)
            rows.append(dict(name=entry['name'],settings=settings,reference=fine,
                refinement_difference={key:abs(fine[key]-coarse[key]) for key in fine if key!='time_days'}))
    args.out.write_text(json.dumps(dict(kind='independent synthetic polar-area and eccentric-anomaly reference, not observations',radial_order=160,angular_order=640,rows=rows),indent=2)+'\n')
    print('AC18 independent polar reference:',len(rows),'cases')


if __name__=='__main__':main()