#!/usr/bin/env python3

import os
import sys
import numpy as np
import phoebe
from astropy import units

from phoebe.parameters import hierarchy as _hierarchy

def default_triple():

    b = phoebe.Bundle()

    b.add_star(component='starA', color='red')
    b.add_star(component='starB', color='green')
    b.add_star(component='starC', color='blue')

    b.add_orbit(component='orbit1', period=1.0*units.d)
    b.add_orbit(component='orbit2', period=10.0*units.d)

    hier1 = _hierarchy.binaryorbit(b['orbit1'], b['starA'], b['starB'])
    hier2 = _hierarchy.binaryorbit(b['orbit2'], hier1, b['starC'])

    b.set_hierarchy(hier2)

    print("hier = ", b.get_hierarchy())

    b.add_constraint(phoebe.parameters.constraint.keplers_third_law_hierarchical, 'orbit2', 'orbit1')
    b.add_compute()

    return b

def test_interferometry5(verbose=True):

    phoebe.conf.devel_on()

    b = default_triple()

    b.set_value('long_an@orbit2@orbit@component', 90.0*units.deg)

    if verbose:
        print("P1 = ", b['period@orbit1@orbit@component'])
        print("P2 = ", b['period@orbit2@orbit@component'])
        print("a1 = ", b['sma@orbit1@orbit@component'])
        print("a2 = ", b['sma@orbit2@orbit@component'])
        print("m1 = ", b['mass@starA@component'])
        print("m2 = ", b['mass@starB@component'])
        print("m3 = ", b['mass@starC@component'])

    sys.exit(1)

    f = open('twigs.txt', 'w')
    for twig in b.twigs:
      f.write("%s\n" % (twig))
    f.close()
 
    b.add_dataset('mesh', times=[2.5+0.25], columns=['visibilities'])
    b.run_compute()
    fig, plt = b.plot(show=True, fc='visibilities', ec=None)
    plt.savefig("mesh.png")
#    sys.exit(1)

    dir_ = os.path.dirname(os.path.realpath(__file__))

    times, u, v, wavelengths, vises, sigmas = np.loadtxt(os.path.join(dir_, "Vis.dat"), usecols=[0, 1, 2, 3, 5, 6], unpack=True)

    b.add_dataset('vis', times=times, u=u, v=v, wavelengths=wavelengths, vises=vises, sigmas=sigmas)

    b.set_value('distance', context = 'system', value=10*units.pc)

    b.set_value('ntriangles@starA', context='compute', value=500)
    b.set_value('ntriangles@starB', context='compute', value=500)
    b.set_value('ntriangles@starC', context='compute', value=500)

    b.run_compute()

    times = b['times@vis01@phoebe01@latest@vis@model'].value
    u = b['u@vis01@phoebe01@latest@vis@model'].value
    v = b['v@vis01@phoebe01@latest@vis@model'].value
    wavelengths = b['wavelengths@vis01@phoebe01@latest@vis@model'].value
    vises = b['vises@vis01@phoebe01@latest@vis@model'].value

    np.savetxt('test_interferometry5.out', np.c_[times, u, v, wavelengths, vises], header='times u v wavelenghts vises')

    #b.plot(show=True)


if __name__ == "__main__":

    #logger = phoebe.logger(clevel='INFO')
    #logger = phoebe.logger(clevel='DEBUG')

    test_interferometry5()


