# -*- coding: utf-8 -*-
"""
Created on Tue Oct 09 14:51:51 2012

@author: Buratti
port 9, port 7 e port 1
"""

import tokdata
import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
from matplotlib.mlab import csd as CSD

def Chan(ns, nchan, t1, t2):
    """Read channel and restrict to t1-t2 range."""
    chname = 'mhdfst.ch{:02}'.format(nchan)
    a = tokdata.ftudata(ns, chname)
    return a.trange([t1, t2]).v, a.ok()

def Phe(ph, th, fi, m, n, th_ref, fi_ref):
    """Distance from expected phase."""
    phe = ph - (m * (th - th_ref) + \
    n * (fi - fi_ref) - 180) % 360 + 180
    phe = (phe + 180) % 360 - 180
    return phe

def mpipi(psi):
    """Put angle in -180:180 deg range"""
    return psi - 360.*np.sign(psi)*np.floor(0.5+np.abs(psi)/360.)

#def plotpsi(m, p, th1, th2, pl, color):
#    kmax = int(np.floor((p + m*(th2-th1))/360. - 0.5))
#    if kmax < 0 :
#        plt.plot([th1, th2], [p, p + m*(th2-th1)], color=color)
#    else:
#        ml = m * pl
#        th = th1 + ((2*np.arange(kmax+1)+1)*180.-p)/m
#        plt.plot([th1, th[0] + pl], [p, 180. + ml], color=color)
#        for k in range(0, kmax):
#            plt.plot([th[k]-pl, th[k+1]+pl], [-180-ml, 180+ml], color=color)
#        plt.plot([th[-1]-pl, th2], [-180-ml, -180+m*(th2-th[-1])], color=color)
def plotpsi(m, p, th1, th2, pl, color):
    sm = np.sign(m)
    am = np.abs(m)
    kmax = np.floor((sm*p + am*(th2-th1))/360. - 0.5).astype(int)
    if kmax < 0 :
        plt.plot([th1, th2], [p, p + m*(th2-th1)], color=color)
    else:
        ml = am * pl
        th = th1 - p/m + ((2*np.arange(kmax+1)+1)*180.)/am
        #print -180-ml, -180+m*(th2-th[-1]), 'y', sm*[-180-ml, -180+m*(th2-th[-1])]
        plt.plot([th1, th[0] + pl], [p, sm*(180. + ml)], color=color)
        for k in range(0, kmax):
            plt.plot([th[k]-pl, th[k+1]+pl], [sm*(-180-ml), sm*(180+ml)], color=color)
        plt.plot([th[-1]-pl, th2], [sm*(-180-ml), -180*sm+m*(th2-th[-1])], color=color)



class Coils(object):
    def __init__(self, ns):
        with open("MHD_pool_all.txt", 'r') as f:
            a = np.genfromtxt(f, names=True)
        # scelta adc in base a ns <<<<<<<<<<<<<
        poolshots = np.array([30959, 32735, 35352, 35405, 35505, \
        35895, 35932, 36041, 36667, 36759])
        config = np.flatnonzero(ns >= poolshots)[-1]
        adc = 'adc' + str(config)
        polarity = 's' + str(config)
        chmask = a[adc] < 0
        self.ch = np.ma.array(a[adc], mask = chmask).compressed().astype(int)
        self.torall = a[adc]
        self.tor = np.ma.array(a['phi'], mask = chmask).compressed()
        polall = a['theta']
        idx = np.flatnonzero(polall < -100.)
        polall[idx] = polall[idx] + 360.
        idx = np.flatnonzero(polall > 260.)
        polall[idx] = polall[idx] - 360.
        self.polall = polall
        self.pol = np.ma.array(polall, mask = chmask).compressed()
        self.s = np.ma.array(a[polarity], mask = chmask).compressed().astype(int)
        self.ik = np.ma.array(a['mark'], mask = chmask).compressed().astype(int)
        self.ns = ns
        
    def cspectra(self, t1, t2, chref=4, fsam=5.e5, sblock=1024, nover=512):
        self.t1 = t1
        self.t2 = t2
        self.chref = chref
        self.ichref = np.flatnonzero(self.ch == chref)[0]
        self.polref = self.pol[self.ichref]
        self.torref = self.tor[self.ichref]
        chanval, changood = Chan(self.ns, chref, t1, t2)
        #vref = self.s[self.ichref] * Chan(self.ns, chref, t1, t2)
        vref = self.s[self.ichref] * chanval
        ll = len(self.ch)
        Auto = np.empty((sblock/2+1, ll))
        Cross = np.empty((sblock/2+1, ll), dtype='c16')
        for k in range(ll):
            chanval, changood = Chan(self.ns, self.ch[k], t1, t2)
            #vchan = self.s[k] * Chan(self.ns, self.ch[k], t1, t2)
            vchan = self.s[k] * chanval
            if changood:
                Auto[:,k], f = CSD(vchan, vchan, sblock, fsam, noverlap=nover)
                Cross[:,k], f = CSD(vref, vchan, sblock, fsam, noverlap=nover)
        self.f = f
        self.Auto = Auto
        self.Cross = Cross
        #autospettri[frequenza, canale]
        #passare sblock, fsam ecc.
        # grafico spettrogramma per scelta tempi
        # passare sblock

    def connection(self, m=2, n=1, probepol = 0., probetor = 90.):
        col = ['y','g',[1,0,1],'r',[1,0.5,0],[0,0.75,1],'b','k']
        q = np.float(m) / n
        x = np.linspace(0.0, m*360.0, 180)
        y = probepol - (x - probetor)  / q
        y = np.mod(y + 100., 360.) - 100.
        idy = y>=250. 
        yy = np.ma.array(y, mask=idy)
        x = np.mod(x, 360.)
        idx = x >= 350.
        xx = np.ma.array(x, mask=idx)
        plt.figure(facecolor='w')
        ha = plt.subplot(1,1,1)
        ha.plot(xx, yy)
        ha.plot(probetor, probepol, marker='*', markersize=14)
        for k, txt in enumerate(self.ch):
            mrk = 'o'
            if self.ch[k] > 16: mrk = 's'
            ha.plot(self.tor[k], self.pol[k], marker=mrk, color=col[self.ik[k]])
            ha.annotate(txt, (self.tor[k], self.pol[k]))
        plt.xlim([-20., 360.])
        plt.xlabel('Toroidal angle', fontsize=16)
        plt.ylabel('Poloidal angle', fontsize=16)
        plt.xticks(fontsize=14)
        plt.yticks(fontsize=14)
        plt.title('#{}   q={}'.format(self.ns, q))
        
    def connection90(self, m=2, n=1, probepol = 0., probetor = 90.):
        col = ['y','g',[1,0,1],'r',[1,0.5,0],[0,0.75,1],'b','k']
        q = np.float(m) / n
        x = np.linspace(0.0, m*360.0, 180)
        y = probepol - (x - probetor)  / q
        y = np.mod(y + 100., 360.) - 100.
        idy = y>=250. 
        yy = np.ma.array(y, mask=idy)
        y90 = y + 90.0
        y90 = np.mod(y90 + 100., 360.) - 100.
        idy90 = y90>=250. 
        yy90 = np.ma.array(y90, mask=idy90)
        x = np.mod(x, 360.)
        idx = x >= 350.
        xx = np.ma.array(x, mask=idx)
        plt.figure(facecolor='w')
        ha = plt.subplot(1,1,1)
        ha.plot(xx, yy)
        ha.plot(xx, yy90, '--')
        ha.plot(probetor, probepol, marker='*', markersize=14)
        for k, txt in enumerate(self.ch):
            mrk = 'o'
            if self.ch[k] > 16: mrk = 's'
            ha.plot(self.tor[k], self.pol[k], marker=mrk, color=col[self.ik[k]])
            ha.annotate(txt, (self.tor[k], self.pol[k]))
        plt.xlim([-20., 360.])
        plt.xlabel('Toroidal angle', fontsize=16)
        plt.ylabel('Poloidal angle', fontsize=16)
        plt.xticks(fontsize=14)
        plt.yticks(fontsize=14)
        plt.title('#{}   q={}'.format(self.ns, q))
        
    def refspc(self, xfreq=False):
        plt.figure(facecolor='w')
        if xfreq:
            plt.plot(1.e-3*self.f, np.log(self.Auto[:, self.ichref]))
            plt.xlabel('Freq (kHz)', fontsize=16)
        else:
            plt.plot(np.log(self.Auto[:, self.ichref]))
            plt.xlabel('spectral bin', fontsize=16)
        plt.ylabel('log spectrum', fontsize=16)
        plt.xticks(fontsize=14)
        plt.yticks(fontsize=14)
        plt.title('#{}   ch{}   t={}-{}s'.format\
        (self.ns, self.chref, self.t1, self.t2))

    def phasemap(self, m=2, n=1, iforced=0):
        #mrk = ['o','o','o','o','o','o','o','o']
        col = ['y','g',[1,0,1],'r',[1,0.5,0],[0,0.75,1],'b','k']
        prol = 10.0
        tors, ind = np.unique(self.tor, return_index=True)
        itop = 2 + np.argmax(self.Auto[2:, self.ichref])
        if iforced > 0: itop = iforced
        aref = np.sqrt(self.Auto[itop, self.ichref])
        phase = np.angle(self.Cross[itop, :], deg=True)
        ampl = np.sqrt(np.abs(self.Auto[itop, :]))
        coher = np.abs(self.Cross[itop, :]) / ampl / aref
        plt.figure(facecolor='w')
        plt.subplots_adjust(hspace=0.001)
        ha = plt.subplot(3,1,1)
        hb = plt.subplot(3,1,2, sharex=ha)
        hc = plt.subplot(3,1,3, sharex=ha)
        for k, txt in enumerate(self.ch):
            mrk = 'o'
            if self.ch[k] > 16: mrk = 's'
            ha.plot(self.pol[k], phase[k], marker=mrk, color=col[self.ik[k]])
            ha.annotate(txt, (self.pol[k], phase[k]))
            hb.plot(self.pol[k], coher[k], marker=mrk, color=col[self.ik[k]])
            hb.annotate(txt, (self.pol[k], coher[k]))
            hc.plot(self.pol[k], ampl[k]/aref, marker=mrk, color=col[self.ik[k]])

        plt.axes(ha)
        plt.title('#{}   t={}-{}s   f={:4.1f} kHz  m={}, n={}'\
        .format(self.ns, self.t1, self.t2, 1.e-3*self.f[itop], m, n))
        plt.ylabel('phase', fontsize=14)
        for k in range(len(tors)):
            idx = np.flatnonzero(self.tor == tors[k])
            th1 = np.amin(self.pol[idx]) - prol
            th2 = np.amax(self.pol[idx]) + prol
            p0 = mpipi(m*(th1-self.polref) + n*(tors[k]-self.torref))
            #print col[self.ik[ind[k]]], p0, th1
            plotpsi(m, p0, th1, th2, prol, col[self.ik[ind[k]]])
        plt.axes(hb)
        plt.ylabel('coherence', fontsize=14)
        plt.ylim([0, 1.1])
        plt.axes(hc)
        plt.ylabel('norm ampl', fontsize=14)
        plt.xlabel('poloidal angle', fontsize=14)
        plt.xticks(fontsize=14)
        plt.show()


    def residuals(self, m=2, n=1, iforced=0):
        #mrk = ['o','o','o','o','o','o','o','o']
        col = ['y','g',[1,0,1],'r',[1,0.5,0],[0,0.75,1],'b','k']
        tors, ind = np.unique(self.tor, return_index=True)
        itop = 2 + np.argmax(self.Auto[2:, self.ichref])
        if iforced > 0: itop = iforced
        phase = np.angle(self.Cross[itop, :], deg=True)
        pdev = Phe(phase, self.pol, self.tor, m, n, self.polref, self.torref)
        plt.figure(facecolor='w')
        for k, txt in enumerate(self.ch):
            mrk = 'o'
            if self.ch[k] > 16: mrk = 's'
            plt.plot(self.pol[k], pdev[k], marker=mrk, color=col[self.ik[k]])
            plt.annotate(txt, (self.pol[k], pdev[k]))
        plt.title('#{}   t={}-{}s   f={:4.1f} kHz  m={}, n={}'\
        .format(self.ns, self.t1, self.t2, 1.e-3*self.f[itop], m, n))
        plt.ylabel('phase deviation', fontsize=14)
        plt.xlabel('poloidal angle', fontsize=14)
        plt.xticks(fontsize=14)
        plt.axhline(y=-50, color='k')
        plt.axhline(y=50, color='k')
        plt.axhline(color='k')
        plt.show()


#ns = 33731; t1 = 0.79; t2 = 0.81 # cycles high dens. NO>17?
#ns = 34769; t1 = 0.69; t2 = 0.7
#ns = 34843; t1 = 0.89; t2 = 0.90 #; iforced = 28 # cycles high dens
#ns = 36252; t1 = 0.60; t2 = 0.62 #
#ns = 37348; t1 = 0.80; t2 = 0.82 # 0.7MA no Ne
#ns = 37548; t1 = 0.53; t2 = 0.55 # IIharm (13); BAE (85, 98)
#ns = 37607; t1 = 0.98; t2 = 1. 
#ns = 38100; t1 = 1.57; t2 = 1.59 #; iforced = 28 # cycles high dens
#ns = 37768; t1 = 1.02; t2 = 1.04 #; iforced = 28 # cycles high dens
#ns = 37768; t1 = 1.02; t2 = 1.04 #; iforced = 28 # cycles high dens
#ns = 38142; t1 = 0.6; t2 = 0.61
#ns = 38142; t1 = 1.146; t2 = 1.152
#ns = 34843; t1 = 0.90; t2 = 0.92
ns = 37655; t1 = 0.33; t2 = 0.35
C = Coils(ns)
C.connection90(probepol=-21, probetor=240)
C.cspectra(t1, t2)
C.phasemap()
#C.residuals(3, 2)
#C.refspc()