
from __future__ import print_function

import sys
import unittest
import math
import numpy
from numpy import asfarray
from numpy import testing as NT
from numpy.testing import assert_array_almost_equal as assert_aequal

from .. import Machine

from .. import GLPSParser

import os
datadir = os.path.dirname(__file__)

def print_state(S):
  n = 7

  sys.stdout.write('\nmoment0_env:\n')
  sys.stdout.write('[')
  for i in range(n):
    sys.stdout.write('%23.16e' % (S.moment0_env[i]))
    if i != n-1: sys.stdout.write(', ')
  sys.stdout.write(']\n')

  sys.stdout.write('moment1_env:\n')
  for i in range(n):
    sys.stdout.write('[')
    for j in range(n):
      sys.stdout.write('%23.16e' % (S.moment1_env[i, j]))
      if j != n-1: sys.stdout.write(', ')
    if i != n-1:
      sys.stdout.write('],\n')
    else:
      sys.stdout.write(']\n')

      sys.stdout.write('IonEk:\n[')
      n1 = len(S.phis)
      for i in range(n1):
        sys.stdout.write('%23.16e' % (S.IonEk[i]))
        if i != n1-1: sys.stdout.write(', ')
      sys.stdout.write(']\n')

      sys.stdout.write('phis:\n[')
      n1 = len(S.phis)
      for i in range(n1):
        sys.stdout.write('%21.14e' % (S.phis[i]))
        if i != n1-1: sys.stdout.write(', ')
      sys.stdout.write(']\n')

class MomentTest(object):
    'Helper for testing moment2 sim'

    def assertConsistent(self, S, msg=None):
        'Check internal consistencies of a State'
        # check "cached" parameters calculated from energy
        self.assertEqual(S.ref_IonW,  S.ref_IonEs+S.ref_IonEk, msg)
        self.assertEqual(S.ref_gamma, S.ref_IonW/S.ref_IonEs, msg)
        self.assertEqual(S.ref_beta,  math.sqrt(1e0-1e0/(S.ref_gamma**2)), msg)
        self.assertEqual(S.ref_bg,    S.ref_beta*S.ref_gamma)

        assert_aequal(S.IonW,  S.IonEs+S.IonEk, err_msg=msg)
        assert_aequal(S.gamma, S.IonW/S.IonEs, err_msg=msg)
        assert_aequal(S.beta,  numpy.sqrt(1e0-1e0/numpy.square(S.gamma)), err_msg=msg)
        assert_aequal(S.bg,    S.beta*S.gamma, err_msg=msg)

        # moment0_env
        W = S.IonQ/S.IonQ.sum()
        W = W.reshape((1,len(W))).repeat(7,axis=0)
        M0env = (S.moment0*W).sum(axis=1)
        assert_aequal(S.moment0_env, M0env, err_msg=msg)

        # moment1_env
        Qs, M0, M0env, M1 = S.IonQ, S.moment0, S.moment0_env, S.moment1
        # using a loop because outer() doesn't understand axis=
        M1env  = numpy.zeros((7,7))
        for i in range(len(Qs)):
            Q = Qs[i]
            m0diff = M0[:,i]-M0env

            M1env[:7,:7] += Q*(M1[...,i]+numpy.outer(m0diff, m0diff))
        M1env /= Qs.sum()

        assert_aequal(S.moment1_env, M1env, err_msg=msg)

        # moment0_rms
        M0rms = numpy.sqrt(numpy.diagonal(S.moment1_env))
        assert_aequal(S.moment0_rms, M0rms, err_msg=msg)

    def assertStateEqual(self, expect, actual, msg=None, decimal=10):
        'Assert that two moment2 State instances are equal'
        for k,v in expect.items():
            #self.assertIn(k, actual, "%s %s missing"%(msg or "",k)) #TODO: make State iterable
            A = getattr(actual, k, None)
            if A is None:
                self.assertTrue(False, "%s %s missing"%(msg or "",k))
            if isinstance(v, numpy.ndarray):
                assert_aequal(v, getattr(actual, k), decimal=decimal, err_msg="%s %s doesn't match"%(msg or "",k))
            else:
                self.assertAlmostEqual(v, getattr(actual, k), places=decimal, msg="%s %s doesn't match"%(msg or "",k))


    def checkPropagate(self, elem, instate, outstate, max=1):
        '''Pass given input state through the named element

        Propagate the same input twice in succession to verify that element(s)
        give the same output.
        The repeat with 'transfer' caching disabled
        '''

        S1 = self.M.allocState(instate, inherit=False)
        S2 = self.M.allocState(instate, inherit=False)

        self.M.propagate(state=S1, start=elem, max=max)
        self.assertConsistent(S1)
        self.M.propagate(state=S2, start=elem, max=max)
        self.assertConsistent(S2)

#        print_state(S1)

        self.assertStateEqual(outstate, S1, 'first pass')
        self.assertStateEqual(outstate, S2, 'second pass') # fails if some Element has a caching problem...

        S1 = self.ICM.allocState(instate, inherit=False)
        S2 = self.ICM.allocState(instate, inherit=False)

        self.ICM.propagate(state=S1, start=elem, max=max)
        self.assertConsistent(S1)
        self.ICM.propagate(state=S2, start=elem, max=max)
        self.assertConsistent(S2)

        self.assertStateEqual(outstate, S1, 'third pass')
        self.assertStateEqual(outstate, S2, 'fourth pass')

class TestToStrl(unittest.TestCase, MomentTest):
    """Strategy is to test the state after the first instance of each element type.

    $ ./tools/flame -N <element_name> \
       -F utest,next_elem,moment0,moment1_env,ref_IonQ,ref_IonZ,ref_IonEs,ref_IonEk,ref_phis,IonQ,IonZ,IonEs,IonEk,phis \
       python/flame/test/to_strl.lat
    """
    lattice = 'to_strl.lat'

    def setUp(self):
        with open(os.path.join(datadir, self.lattice), 'rb') as F:
            self.M = Machine(F)
            F.seek(0)
            self.ICM = Machine(F, extra={'skipcache':1.0})

    def test_source(self):
        # S
        self.checkPropagate(0, {}, {
            'next_elem':1,
            'moment0':asfarray([
                [-7.8859999999999998e-04, 7.3126670000000001e-03],
                [1.0837100000000000e-05, 1.4765650000000000e-05],
                [1.3373430000000000e-02, 3.4498490000000001e-03],
                [6.6785340000000002e-06, -7.3976846899999996e-06],
                [-1.8477289999999999e-04, 2.2978860000000000e-02],
                [3.0999499999999998e-04, 2.0601000000000000e-03],
                [1.0000000000000000e+00, 1.0000000000000000e+00],
            ]),
            'moment1_env':asfarray([
                [2.7784895409858517e+00, 2.4291087928314132e-04, -2.8819146985954898e-03, -1.3098890046625701e-05, 5.3594209519492754e-04, -1.7328482801286655e-05, 0.0000000000000000e+00],
                [2.4291087928314132e-04, 4.0675140732215713e-06, 9.1913901431043622e-06, -2.2427664519913498e-08, 3.5190231444234908e-07, -1.1570581361000062e-10, 0.0000000000000000e+00],
                [-2.8819146985954898e-03, 9.1913901431043622e-06, 2.5433968050163540e+00, 2.3922914294383500e-04, -8.7109112520915007e-04, -2.3018796435535812e-06, 0.0000000000000000e+00],
                [-1.3098890046625701e-05, -2.2427664519913498e-08, 2.3922914294383500e-04, 4.6600708350882344e-06, -8.7385463545847361e-07, 1.7247530507817648e-08, 0.0000000000000000e+00],
                [5.3594209519492754e-04, 3.5190231444234908e-07, -8.7109112520915007e-04, -8.7385463545847361e-07, 1.0267518939868806e-03, 1.0056757657182434e-05, 0.0000000000000000e+00],
                [-1.7328482801286655e-05, -1.1570581361000062e-10, -2.3018796435535812e-06, 1.7247530507817648e-08, 1.0056757657182434e-05, 2.6314343481397998e-06, 0.0000000000000000e+00],
                [0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00],
            ]),
            'ref_IonQ':1.01110000000000000e+04,
            'ref_IonZ':1.38655462184873957e-01,
            'ref_IonEs':9.31494320000000000e+08,
            'ref_IonEk':5.00000000000000000e+05,
            'ref_phis':0.00000000000000000e+00,
            'IonQ':asfarray([1.0111000000000000e+04, 1.0531000000000000e+04]),
            'IonZ':asfarray([1.3865546218487396e-01, 1.4285714285714285e-01]),
            'IonEs':asfarray([9.3149432000000000e+08, 9.3149432000000000e+08]),
            'IonEk':asfarray([5.0030999500000000e+05, 5.0206009999999998e+05]),
            'phis':asfarray([-1.8477289999999999e-04, 2.2978860000000000e-02]),
        }, max=1)

    def test_drift(self):
        # drift_1
        self.checkPropagate(0, {}, {
            'next_elem':2,
            'moment0':asfarray([
                [-8.3287999999999938e-06, 8.3757937999999997e-03],
                [1.0837100000000000e-05, 1.4765650000000000e-05],
                [1.3854284448000000e-02, 2.9172157023199998e-03],
                [6.6785340000000002e-06, -7.3976846899999996e-06],
                [-1.3325396829281771e-03, 1.5391130160821812e-02],
                [3.0999499999999998e-04, 2.0601000000000000e-03],
                [1.0000000000000000e+00, 1.0000000000000000e+00],
            ]),
            'moment1_env':asfarray([
                [2.8345547005582050e+00, 5.3577189255509439e-04, -3.2795197045202571e-03, -1.4713681892059471e-05, 6.2542645808289041e-04, -1.7336813619866573e-05, 0.0000000000000000e+00],
                [5.3577189255509439e-04, 4.0675140732215713e-06, 7.5765982976705889e-06, -2.2427664519913498e-08, 3.5228427569447906e-07, -1.1570581361000062e-10, 0.0000000000000000e+00],
                [-3.2795197045202571e-03, 7.5765982976705889e-06, 2.6020036088093637e+00, 5.7475424307018780e-04, -9.3022235763712518e-04, -1.0600574469907109e-06, 0.0000000000000000e+00],
                [-1.4713681892059471e-05, -2.2427664519913498e-08, 5.7475424307018780e-04, 4.6600708350882344e-06, -9.3792991959628860e-07, 1.7247530507817648e-08, 0.0000000000000000e+00],
                [6.2542645808289041e-04, 3.5228427569447906e-07, -9.3022235763712475e-04, -9.3792991959628860e-07, 9.8879108725721874e-04, 3.4841141433569607e-07, 0.0000000000000000e+00],
                [-1.7336813619866573e-05, -1.1570581361000062e-10, -1.0600574469907109e-06, 1.7247530507817648e-08, 3.4841141433569607e-07, 2.6314343481397998e-06, 0.0000000000000000e+00],
                [0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00],
            ]),
            'ref_IonQ':1.01110000000000000e+04,
            'ref_IonZ':1.38655462184873957e-01,
            'ref_IonEs':9.31494320000000000e+08,
            'ref_IonEk':5.00000000000000000e+05,
            'ref_phis':3.70896240162486901e+00,
            'IonQ':asfarray([1.0111000000000000e+04, 1.0531000000000000e+04]),
            'IonZ':asfarray([1.3865546218487396e-01, 1.4285714285714285e-01]),
            'IonEs':asfarray([9.3149432000000000e+08, 9.3149432000000000e+08]),
            'IonEk':asfarray([5.0030999500000000e+05, 5.0206009999999998e+05]),
            'phis':asfarray([3.7076293281508645e+00, 3.7243300943218349e+00]),
        }, max=2)

    def test_rfcav_41(self):
        # ls1_ca01_cav1_d1127  cavtype = "0.041QWR"
        self.assertEqual(self.M.find(name='ls1_ca01_cav1_d1127')[0], 3)
        self.checkPropagate(0, {}, {
            'next_elem':4,
            'moment0':asfarray([
                [4.2806343239717116e-03, 1.5129972255810526e-02],
                [1.3170791010251590e-05, 2.6638715098587717e-05],
                [4.0306772693611600e-03, -1.4238677753316846e-02],
                [-1.6893666125019258e-04, -2.0313868893374049e-04],
                [-5.8199847180944175e-03, -2.8171535425187244e-02],
                [1.0937674260139465e-04, 3.2640186251401901e-03],
                [1.0000000000000000e+00, 1.0000000000000000e+00],
            ]),
            'moment1_env':asfarray([
                [4.5737889068082804e+00, 6.0669648277006495e-03, -1.0792259945136952e-02, -3.1759168650761263e-05, 9.1002516920943549e-04, 3.0358922802964227e-05, 0.0000000000000000e+00],
                [6.0669648277006495e-03, 1.0279529425767875e-05, -1.0736164288211990e-05, -4.6551002950065074e-08, 1.0542984435954559e-06, 4.2234819235438213e-08, 0.0000000000000000e+00],
                [-1.0792259945136953e-02, -1.0736164288211994e-05, 4.4609155106701239e+00, 6.1807101898400177e-03, -1.1255393909196713e-03, -6.3650656953149896e-05, 0.0000000000000000e+00],
                [-3.1759168650761263e-05, -4.6551002950065080e-08, 6.1807101898400177e-03, 1.0961176841497697e-05, -1.8057731753519633e-06, -9.3605511068373278e-08, 0.0000000000000000e+00],
                [9.1002516920943570e-04, 1.0542984435954559e-06, -1.1255393909196708e-03, -1.8057731753519637e-06, 1.2760917189525325e-03, -4.6806127720643879e-08, 0.0000000000000000e+00],
                [3.0358922802964227e-05, 4.2234819235438213e-08, -6.3650656953149896e-05, -9.3605511068373278e-08, -4.6806127720648921e-08, 4.2230866648902030e-06, 0.0000000000000000e+00],
                [0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00],
            ]),
            'ref_IonQ':1.01110000000000000e+04,
            'ref_IonZ':1.38655462184873957e-01,
            'ref_IonEs':9.31494320000000000e+08,
            'ref_IonEk':5.53528799277067184e+05,
            'ref_phis':2.27471719868534592e+01,
            'IonQ':asfarray([1.0111000000000000e+04, 1.0531000000000000e+04]),
            'IonZ':asfarray([1.3865546218487396e-01, 1.4285714285714285e-01]),
            'IonEs':asfarray([9.3149432000000000e+08, 9.3149432000000000e+08]),
            'IonEk':asfarray([5.5363817601966858e+05, 5.5679281790220737e+05]),
            'phis':asfarray([2.2741352002135365e+01, 2.2719000451428272e+01]),
        }, max=4)

    def test_rfcav_41_recalcttf(self):
        # ls1_ca01_cav1_d1127  cavtype = "0.041QWR"
        self.assertEqual(self.M.find(name='ls1_ca01_cav1_d1127')[0], 3)
        self.M.reconfigure(3, {'forcettfcalc':1.0})
        self.ICM.reconfigure(3, {'forcettfcalc':1.0})
        self.checkPropagate(0, {}, {
            'next_elem':4,
            'moment0_env':asfarray(
                [ 9.81245052404105e-03,  2.00474944555705e-05, -4.85925056656807e-03, -1.83712372016754e-04, -1.72231520708894e-02,  1.71879122109263e-03,  1.00000000000000e+00]
            ),
            'moment1_env':asfarray([
                [ 4.56978111860199e+00,  6.06523462170448e-03, -1.07862302961237e-02, -3.17725595594832e-05,  9.09605011212125e-04,  3.03125969710196e-05,  0.00000000000000e+00],
                [ 6.06523462170448e-03,  1.02839738079863e-05, -1.07352561610397e-05, -4.65886065788557e-08,  1.05430063236264e-06,  4.21997654070090e-08,  0.00000000000000e+00],
                [-1.07862302961237e-02, -1.07352561610397e-05,  4.45894182337220e+00,  6.18816091797052e-03, -1.12533711708629e-03, -6.35830908045431e-05,  0.00000000000000e+00],
                [-3.17725595594832e-05, -4.65886065788557e-08,  6.18816091797052e-03,  1.09867452299131e-05, -1.80852275365200e-06, -9.35806056450271e-08,  0.00000000000000e+00],
                [ 9.09605011212125e-04,  1.05430063236264e-06, -1.12533711708629e-03, -1.80852275365200e-06,  1.27614016515051e-03, -8.88103508147561e-08,  0.00000000000000e+00],
                [ 3.03125969710196e-05,  4.21997654070090e-08, -6.35830908045431e-05, -9.35806056450271e-08, -8.88103508147579e-08,  4.22170469127191e-06,  0.00000000000000e+00],
                [ 0.00000000000000e+00,  0.00000000000000e+00,  0.00000000000000e+00,  0.00000000000000e+00,  0.00000000000000e+00,  0.00000000000000e+00,  0.00000000000000e+00]
            ]),
            'ref_IonQ':1.01110000000000000e+04,
            'ref_IonZ':1.38655462184873957e-01,
            'ref_IonEs':9.31494320000000000e+08,
            'ref_IonEk':5.53528799277067184e+05,
            'ref_phis':2.27471719868534592e+01,
            'IonQ':asfarray([1.0111000000000000e+04, 1.0531000000000000e+04]),
            'IonZ':asfarray([1.3865546218487396e-01, 1.4285714285714285e-01]),
            'IonEs':asfarray([9.3149432000000000e+08, 9.3149432000000000e+08]),
            'IonEk':asfarray([5.5363817601966858e+05, 5.5679281790220737e+05]),
            'phis':asfarray([2.2741352002135365e+01, 2.2719000451428272e+01]),
        }, max=4)

    def test_solenoid(self):
        # ls1_ca01_sol1_d1131_1
        self.checkPropagate(0, {}, {
            'next_elem':9,
            'moment0':asfarray([
                [-7.5976405161956887e-03, -4.2033140434120323e-03],
                [-3.8191576548473962e-05, -4.3132441823442950e-05],
                [-4.3990633733480659e-02, -7.3762251080284197e-02],
                [-1.1896171910899926e-04, -1.1594740503684112e-04],
                [-7.2514286712388474e-03, -7.0526029188607994e-02],
                [1.0937674260139465e-04, 3.2640186251401901e-03],
                [1.0000000000000000e+00, 1.0000000000000000e+00],
            ]),
            'moment1_env':asfarray([
                [8.0532835805828107e+00, -1.1996773374006241e-03, -1.2076556960799312e-02, 7.8260918480707844e-05, 4.4722636729343661e-04, 7.7455375507305819e-06, 0.0000000000000000e+00],
                [-1.1996773374006249e-03, 1.4386139131164235e-06, 9.9286883345759832e-05, -3.4739396768956866e-09, -2.5932535513055464e-07, -3.8185010164746997e-09, 0.0000000000000000e+00],
                [-1.2076556960799720e-02, 9.9286883345759859e-05, 8.0927099123362876e+00, -9.7902740600577039e-04, -6.4020082997541831e-04, -9.4677007214414072e-05, 0.0000000000000000e+00],
                [7.8260918480708129e-05, -3.4739396768954240e-09, -9.7902740600577104e-04, 1.4431734180561974e-06, -4.0226350811848745e-08, 1.6715286131978529e-08, 0.0000000000000000e+00],
                [4.4722636729343661e-04, -2.5932535513055506e-07, -6.4020082997541809e-04, -4.0226350811849354e-08, 1.9877692143781063e-03, -5.4891187320312006e-05, 0.0000000000000000e+00],
                [7.7455375507305819e-06, -3.8185010164747047e-09, -9.4677007214414072e-05, 1.6715286131978532e-08, -5.4891187320312006e-05, 4.2230866648902030e-06, 0.0000000000000000e+00],
                [0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00],
            ]),
            'ref_IonQ':1.01110000000000000e+04,
            'ref_IonZ':1.38655462184873957e-01,
            'ref_IonEs':9.31494320000000000e+08,
            'ref_IonEk':5.53528799277067184e+05,
            'ref_phis':3.72527602222924656e+01,
            'IonQ':asfarray([1.0111000000000000e+04, 1.0531000000000000e+04]),
            'IonZ':asfarray([1.3865546218487396e-01, 1.4285714285714285e-01]),
            'IonEs':asfarray([9.3149432000000000e+08, 9.3149432000000000e+08]),
            'IonEk':asfarray([5.5363817601966858e+05, 5.5679281790220737e+05]),
            'phis':asfarray([3.7245508581428005e+01, 3.7182047005562815e+01]),
        }, max=9)

    def test_orbtrim(self):
        # ls1_ca01_dch_d1131_2
        self.checkPropagate(0, {}, {
            'next_elem':10,
            'moment0':asfarray([
                [-7.5976405161956887e-03, -4.2033140434120323e-03],
                [-3.8191576548473962e-05, -4.3132441823442950e-05],
                [-4.3990633733480659e-02, -7.3762251080284197e-02],
                [-1.1896171910899926e-04, -1.1594740503684112e-04],
                [-7.2514286712388474e-03, -7.0526029188607994e-02],
                [1.0937674260139465e-04, 3.2640186251401901e-03],
                [1.0000000000000000e+00, 1.0000000000000000e+00],
            ]),
            'moment1_env':asfarray([
                [8.0532835805828107e+00, -1.1996773374006241e-03, -1.2076556960799312e-02, 7.8260918480707844e-05, 4.4722636729343661e-04, 7.7455375507305819e-06, 0.0000000000000000e+00],
                [-1.1996773374006249e-03, 1.4386139131164235e-06, 9.9286883345759832e-05, -3.4739396768956866e-09, -2.5932535513055464e-07, -3.8185010164746997e-09, 0.0000000000000000e+00],
                [-1.2076556960799720e-02, 9.9286883345759859e-05, 8.0927099123362876e+00, -9.7902740600577039e-04, -6.4020082997541831e-04, -9.4677007214414072e-05, 0.0000000000000000e+00],
                [7.8260918480708129e-05, -3.4739396768954240e-09, -9.7902740600577104e-04, 1.4431734180561974e-06, -4.0226350811848745e-08, 1.6715286131978529e-08, 0.0000000000000000e+00],
                [4.4722636729343661e-04, -2.5932535513055506e-07, -6.4020082997541809e-04, -4.0226350811849354e-08, 1.9877692143781063e-03, -5.4891187320312006e-05, 0.0000000000000000e+00],
                [7.7455375507305819e-06, -3.8185010164747047e-09, -9.4677007214414072e-05, 1.6715286131978532e-08, -5.4891187320312006e-05, 4.2230866648902030e-06, 0.0000000000000000e+00],
                [0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00, 0.0000000000000000e+00],
            ]),
            'ref_IonQ':1.01110000000000000e+04,
            'ref_IonZ':1.38655462184873957e-01,
            'ref_IonEs':9.31494320000000000e+08,
            'ref_IonEk':5.53528799277067184e+05,
            'ref_phis':3.72527602222924656e+01,
            'IonQ':asfarray([1.0111000000000000e+04, 1.0531000000000000e+04]),
            'IonZ':asfarray([1.3865546218487396e-01, 1.4285714285714285e-01]),
            'IonEs':asfarray([9.3149432000000000e+08, 9.3149432000000000e+08]),
            'IonEk':asfarray([5.5363817601966858e+05, 5.5679281790220737e+05]),