# -*- coding: utf-8 -*-
"""
Copyright or © or Copr. Project-Team S3 (Sound Signals and Systems) and Analysis/Synthesis team,
Laboratory of Sciences and Technologies of Music and Sound (UMR 9912),
IRCAM-CNRS-UPMC, 1 place Igor Stravinsky, F-75004 Paris
contributor(s) : Antoine Falaize, Thomas Hélie, Thu Jul 9 23:11:37 2015
corresponding contributor: antoine.falaize@ircam.fr

This software (pypHs) is a computer program whose purpose is to generate C++ 
code for the simulation of multiphysics system described in graph formalism. 
It is composed of a library (pypHs.py) and a dictionnary (Dictionnary.py)

This software is governed by the CeCILL-B license under French law and
abiding by the rules of distribution of free software.  You can  use, 
modify and/ or redistribute the software under the terms of the CeCILL-B
license as circulated by CEA, CNRS and INRIA at the following URL
"http://www.cecill.info". 

As a counterpart to the access to the source code and  rights to copy,
modify and redistribute granted by the license, users are provided only
with a limited warranty  and the software's author,  the holder of the
economic rights,  and the successive licensors  have only  limited
liability. 

In this respect, the user's attention is drawn to the risks associated
with loading,  using,  modifying and/or developing or reproducing the
software by the user in light of its specific status of free software,
that may mean  that it is complicated to manipulate,  and  that  also
therefore means  that it is reserved for developers  and  experienced
professionals having in-depth computer knowledge. Users are therefore
encouraged to load and test the software's suitability as regards their
requirements in conditions enabling the security of their systems and/or 
data to be ensured and,  more generally, to use and operate it in the 
same conditions as regards security. 

The fact that you are presently reading this means that you have had
knowledge of the CeCILL-B license and that you accept its terms.

"""

from pypHs import PortHamiltonianObject
import sympy as sp
import numpy as np
import networkx as nx
from sympy.physics import units


### Electrical domain ###

class capacitor(PortHamiltonianObject):
    """ Linear Capacitor
    usgae: capacitor label ['n1','n2'] [value]

    """
    def __init__(self, label, nodes_labels, value):
        
        PortHamiltonianObject.__init__(self,label)
        if type(value[0])==tuple:
            string = value[0][0]
            C = sp.symbols(string)
            subs = [(C, value[0][1])]
            pars = [C]
        else:
            C = value[0]
            subs = []            
            pars = []
        self.subs_par += subs
        self.params += pars

        self.AddLinearStorageComponents(["x"+label],[1./C], [units.C])
        edge_data_dic = {'ref':"x"+label, 'type':'storage', 'realizability':'flux_controlled', 'linear':True, 'link_ref':"x"+label}
        self.Graph.add_edges_from([(nodes_labels[0], nodes_labels[1], edge_data_dic)])

class inductor(PortHamiltonianObject):
    """ Linear Inductor
    """
    def __init__(self,label,nodes_labels,value):

        if type(value[0])==tuple:
            string = value[0][0]
            L = sp.symbols(string)
            subs = [(L, value[0][1])]
            pars = [L]
        else:
            L = value[0]
            subs = []            
            pars = []

        PortHamiltonianObject.__init__(self,label)
        self.AddLinearStorageComponents(["x"+label],[1./L], [units.magnetic_flux])
        edge_data_dic = {'ref':"x"+label, 'type':'storage', 'realizability':'effort_controlled', 'linear':True, 'link_ref':"x"+label}
        self.Graph.add_edges_from([(nodes_labels[0], nodes_labels[1], edge_data_dic)])
        self.subs_par += subs
        self.params += pars

class resistor(PortHamiltonianObject):
    """ Linear Resistor
    """
    def __init__(self,label,nodes_labels,value):
        if type(value[0])==tuple:
            string = value[0][0]
            R = sp.symbols(string)
            pars = [R]
            subs = [(R, value[0][1])]
        else:
            R = value[0]
            subs = []
            pars = []
        PortHamiltonianObject.__init__(self,label)
        self.AddLinearDissipativeComponents(["w"+label],[R], [units.A])
        edge_data_dic = {'ref':"w"+label, 'type':'dissipative', 'realizability':'?', 'linear':True, 'link_ref':"w"+label}
        self.Graph.add_edges_from([(nodes_labels[0], nodes_labels[1], edge_data_dic)])
        self.subs_par += subs
        self.params += pars

class pn_diode(PortHamiltonianObject):
    """ Linear Resistor
    """
    def __init__(self,label,nodes_labels,value):
        PortHamiltonianObject.__init__(self,label)
        subs = []
        pars = []
        if type(value[0])==tuple:
            string = value[0][0]
            Is = sp.symbols(string)
            subs += [(Is, value[0][1])]
            pars += [Is]
        else:
            Is = value[0]
        if type(value[1])==tuple:
            string = value[1][0]
            v0 = sp.symbols(string)
            subs += [(v0, value[1][1])]
            pars += [v0]
        else:
            v0 = value[1]
        self.subs_par += subs
        self.params += pars
        w = sp.symbols("w"+label)
        expr = Is*(sp.exp(w/v0)-1)
        self.AddNonLinearDissipativeComponents([w],[expr], [units.V])
        edge_data_dic = {'ref':str(w), 'type':'dissipative', 'realizability':'effort_controlled', 'linear':False, 'link_ref':str(w)}
        self.Graph.add_edges_from([(nodes_labels[0], nodes_labels[1], edge_data_dic)])

class port(PortHamiltonianObject):
    """ port
    """
    def __init__(self,label,nodes_labels,value):
        PortHamiltonianObject.__init__(self,label)
        u, y = sp.symbols('u_'+label+', y_'+label)        
        unit = units.V if value[0]=='f' else units.A
        self.AddPorts([u],[y],[unit])
        if value.__len__()>1:
            self.AddConstantInpu('u_'+label, value[1])
        realizability = "effort_controlled" if value[0]=='f' else "flux_controlled"
        edge_data_dic = {'ref':str(y), 'type':'port', 'realizability':realizability, 'linear':'?', 'link_ref':str(y)}
        self.Graph.add_edges_from([(0, nodes_labels[0], edge_data_dic)])

class gyrator(PortHamiltonianObject):
    """ 
    gyrator label [N1,N2,N3,N4] [(str, float)]

    consider e1 = value*f2; e2 = value*f1 with apprpriate units  

    """
    def __init__(self,label,nodes_labels,value):
        PortHamiltonianObject.__init__(self,label)
        if type(value[0])==tuple:
            string = value[0][0]
            symb = sp.symbols(string)
            subs = [(symb, value[0][1])]
        else:
            string = value[0]
            subs = []            
        self.subs_par += subs
        self.AddTransformers(label, "gyrator", string)        
        labels = ['y_'+label+'_'+str(e) for e in (1,2)]
        edge_data_dic1 = {'type':'trans_port', 'realizability':'?', 'ref':labels[0], 'link_ref':labels[1], 'linear':'?'}
        edge_data_dic2 = {'type':'trans_port', 'realizability':'?', 'ref':labels[1], 'link_ref':labels[0], 'linear':'?'}
        N1, N2, N3, N4 = nodes_labels
        edges = [(N1, N2, edge_data_dic1),(N3, N4, edge_data_dic2)]
        self.Graph.add_edges_from(edges)

class transformer(PortHamiltonianObject):
    """ 
    gyrator label [N1,N2,N3,N4] [value, unit]

    consider e1 = value*f2; e2 = value*f1 with apprpriate units  

    """
    def __init__(self,label,nodes_labels,value):
        PortHamiltonianObject.__init__(self,label)
        if type(value[0])==tuple:
            string = value[0][0]
            symb = sp.symbols(string)
            subs = [(symb, value[0][1])]
        else:
            string = value[0]
            subs = []            
        self.subs_par += subs
        self.AddTransformers(label, "transformer", string)
        labels = ['y_'+label+'_'+str(e) for e in (1,2)]
        edge_data_dic1 = {'type':'trans_port', 'realizability':'?', 'ref':labels[0], 'link_ref':labels[1], 'linear':'?'}
        edge_data_dic2 = {'type':'trans_port', 'realizability':'?', 'ref':labels[1], 'link_ref':labels[0], 'linear':'?'}
        N1, N2, N3, N4 = nodes_labels
        edges = [(N1, N2, edge_data_dic1),(N3, N4, edge_data_dic2)]
        self.Graph.add_edges_from(edges)

# Mechanical domain

class stiffness(PortHamiltonianObject):
    """ Linear Capacitor
    usgae: capacitor label ['n1','n2'] [value]

    """
    def __init__(self, label, nodes_labels, value):
        PortHamiltonianObject.__init__(self,label)
        if type(value[0])==tuple:
            string = value[0][0]
            K = sp.symbols(string)
            subs = [(K, value[0][1])]
            pars = [K]
        else:
            K = value[0]
            subs = []            
            pars = []
        self.subs_par += subs
        self.params += pars

        self.AddLinearStorageComponents(["x"+label],[K], [units.m])
        edge_data_dic = {'ref':"x"+label, 'type':'storage', 'realizability':'flux_controlled', 'linear':True, 'link_ref':"x"+label}
        self.Graph.add_edges_from([(nodes_labels[0], nodes_labels[1], edge_data_dic)])
        
class mass(PortHamiltonianObject):
    """ Linear Inductor
    """
    def __init__(self,label,nodes_labels,value):
        PortHamiltonianObject.__init__(self,label)
        if type(value[0])==tuple:
            string = value[0][0]
            m = sp.symbols(string)
            subs = [(m, value[0][1])]
            pars = [m]
        else:
            m = value[0]
            subs = []            
            pars = []
        self.subs_par += subs
        self.params += pars

        self.AddLinearStorageComponents(["x"+label],[1./m], [units.kg*units.m/units.s])
        edge_data_dic = {'ref':"x"+label, 'type':'storage', 'realizability':'effort_controlled', 'linear':True, 'link_ref':"x"+label}
        self.Graph.add_edges_from([(nodes_labels[0], nodes_labels[1], edge_data_dic)])

class damper(PortHamiltonianObject):
    """ Linear Resistor
    """
    def __init__(self,label,nodes_labels,value):
        PortHamiltonianObject.__init__(self,label)
        if type(value[0])==tuple:
            string = value[0][0]
            a = sp.symbols(string)
            subs = [(a, value[0][1])]
            pars = [a]
        else:
            a = value[0]
            subs = []            
            pars = []
        self.subs_par += subs
        self.params += pars

        self.AddLinearDissipativeComponents(["w"+label],[a], [units.m/units.s])
        edge_data_dic = {'ref':"w"+label, 'type':'dissipative', 'realizability':'?', 'linear':True, 'link_ref':"w"+label}
        self.Graph.add_edges_from([(nodes_labels[0], nodes_labels[1], edge_data_dic)])

class NLstiffness(PortHamiltonianObject):
    """ Linear Capacitor
    usgae: capacitor label ['n1','n2'] [K0,ParametreNL,xsat]

    """
    def __init__(self, label, nodes_labels, value):
        PortHamiltonianObject.__init__(self,label)
        subs = []
        pars = []
        if type(value[0])==tuple:
            string = value[0][0]
            K0 = sp.symbols(string)
            subs += [(K0, value[0][1])]
            pars += [K0]
        else:
            K0 = value[0]
        if type(value[1])==tuple:
            string = value[1][0]
            ParametreNL = sp.symbols(string)
            subs += [(ParametreNL, value[1][1])]
            pars += [ParametreNL]
        else:
            ParametreNL = value[1]
        if type(value[2])==tuple:
            string = value[1][0]
            xsat = sp.symbols(string)
            subs += [(xsat, value[1][1])]
            pars += [xsat]
        else:
            xsat = value[1]
        self.subs_par += subs
        self.params += pars
        x = sp.symbols("x"+label)
        StateSaturatingEnergy = lambda x,c,par,xsat: c*(x**2/2-(8*par*xsat)*( sp.log((sp.cos(((sp.pi*x)/(2*xsat))))) + 0.5*((sp.pi*x)/(2*xsat))**2 )/(sp.pi*(4-sp.pi)) )
        Hnl = StateSaturatingEnergy(x,K0,ParametreNL,xsat)
        self.AddNonLinearStorageComponents([x],Hnl, [units.m])
        edge_data_dic = {'ref':"x"+label, 'type':'storage', 'realizability':'flux_controlled', 'linear':False, 'link_ref':"x"+label}
        self.Graph.add_edges_from([(nodes_labels[0], nodes_labels[1], edge_data_dic)])
        
class FracDerEffortCtrl(PortHamiltonianObject):
    """ Fractional Flux Controlled storage element
    usgae: FracIntEffortCtrl label ['n1','n2'] [rAlphaMag, alphaMag ,NbPoles]

    """
    def __init__(self, label, nodes_labels, value):
        PortHamiltonianObject.__init__(self,label)
        diagRmu, diagQmu = ComputeFractionalDerivativeWeights(value[0], value[1], NbPoles=value[2], NbFreqPoints=200, DoPlot=True,PolesMinMax = (-5,5))
        nbPoles = diagRmu.__len__()
        subs = []
        for n in range(nbPoles):
            self.AddLinearDissipativeComponents(['w'+label+str(n)],[diagRmu[n]], [1])
            edge_data_dic = {'ref':'w'+label+str(n), 'type':'dissipative', 'realizability':'effort_controlled', 'linear':True, 'link_ref':'w'+label+str(n)}
            self.Graph.add_edges_from([(nodes_labels[0], 'N'+label+str(n)+"_1", edge_data_dic)])

            self.AddLinearStorageComponents(["x"+label+str(n)],[diagQmu[n]], [1])
            edge_data_dic = {'ref':"x"+label+str(n), 'type':'storage', 'realizability':'flux_controlled', 'linear':True, 'link_ref':"x"+label+str(n)}
            self.Graph.add_edges_from([('N'+label+str(n)+"_2", 0, edge_data_dic)])

            Symb = sp.symbols('alpha'+label+str(n))
            self.AddTransformers(label+str(n), "transformer", Symb)        
            labels = ['y_'+label+str(n)+'_'+str(e) for e in (1,2)]
            edge_data_dic1 = {'type':'trans_port', 'realizability':'?', 'ref':labels[0], 'link_ref':labels[1], 'linear':'?'}
            edge_data_dic2 = {'type':'trans_port', 'realizability':'?', 'ref':labels[1], 'link_ref':labels[0], 'linear':'?'}
            edges = [('N'+label+str(n)+"_1", nodes_labels[1], edge_data_dic1),('N'+label+str(n)+"_2", 0, edge_data_dic2)]
            self.Graph.add_edges_from(edges)

            subs += [(Symb, diagRmu[n]**-1)]

        self.subs_par += subs

class FracDerFluxCtrl(PortHamiltonianObject):
    """ Fractional Flux Controlled storage element
    usgae: FracIntEffortCtrl label ['n1','n2'] [rAlphaMag, alphaMag ,NbPoles]

    """
    def __init__(self, label, nodes_labels, value):
        PortHamiltonianObject.__init__(self,label)
        diagRmu, diagQmu = ComputeFractionalDerivativeWeights(value[0], value[1], NbPoles=value[2], NbFreqPoints=200, DoPlot=True,PolesMinMax = (-5,5))
        nbPoles = diagRmu.__len__()
        subs = []
        nodes = [nodes_labels[0],]+['N'+label+str(n) for n in range(1,nbPoles)] + [nodes_labels[1],]
        nodes.reverse()
        for n in range(nbPoles):
            N1 = nodes.pop()
            N2 = nodes[-1]

            self.AddLinearDissipativeComponents(['w'+label+str(n)],[diagRmu[n]], [1])
            edge_data_dic = {'ref':'w'+label+str(n), 'type':'dissipative', 'realizability':'flux_controlled', 'linear':True, 'link_ref':'w'+label+str(n)}
            self.Graph.add_edges_from([(N1, N2, edge_data_dic)])

            self.AddLinearStorageComponents(["x"+label+str(n)],[diagQmu[n]], [1])
            edge_data_dic = {'ref':"x"+label+str(n), 'type':'storage', 'realizability':'effort_controlled', 'linear':True, 'link_ref':"x"+label+str(n)}
            self.Graph.add_edges_from([(str(N1)+"trans"+str(n), 0, edge_data_dic)])

            Symb = sp.symbols('alpha'+label+str(n))
            self.AddTransformers(label+str(n), "transformer", Symb)        
            labels = ['y_'+label+str(n)+'_'+str(e) for e in (1,2)]
            edge_data_dic1 = {'type':'trans_port', 'realizability':'?', 'ref':labels[0], 'link_ref':labels[1], 'linear':'?'}
            edge_data_dic2 = {'type':'trans_port', 'realizability':'?', 'ref':labels[1], 'link_ref':labels[0], 'linear':'?'}
            edges = [(N1, N2, edge_data_dic1),(str(N1)+"trans"+str(n), 0, edge_data_dic2)]
            self.Graph.add_edges_from(edges)

            subs += [(Symb, diagRmu[n]**-1)]

        self.subs_par += subs

class FracIntEffortCtrl(PortHamiltonianObject):
    """ Fractional Flux Controlled storage element
    usgae: FracIntEffortCtrl label ['n1','n2'] [rAlphaMag, beta ,NbPoles]

    """
    def __init__(self, label, nodes_labels, value):
        PortHamiltonianObject.__init__(self,label)
        diagRmu, diagQmu = ComputeFractionalDerivativeWeights(value[0], value[1], NbPoles=value[2], NbFreqPoints=200, DoPlot=True,PolesMinMax = (-5,5),out_integrator = True)
        nbPoles = diagRmu.__len__()
        for n in range(nbPoles):
            self.AddLinearDissipativeComponents(['w'+label+str(n)],[diagRmu[n]], [1])
            edge_data_dic = {'ref':'w'+label+str(n), 'type':'dissipative', 'realizability':'flux_controlled', 'linear':True, 'link_ref':'w'+label+str(n)}
            self.Graph.add_edges_from([('node'+label+str(n),nodes_labels[1], edge_data_dic)])

            self.AddLinearStorageComponents(["x"+label+str(n)],[diagQmu[n]], [1])
            edge_data_dic = {'ref':"x"+label+str(n), 'type':'storage', 'realizability':'effort_controlled', 'linear':True, 'link_ref':"x"+label+str(n)}
            self.Graph.add_edges_from([(nodes_labels[0],'node'+label+str(n), edge_data_dic)])

class FracIntFluxCtrl(PortHamiltonianObject):
    """ Fractional Flux Controlled storage element
    usgae: FracIntFluxCtrl label ['n1','n2'] [rAlphaMag, alphaMag ,NbPoles]

    """
    def __init__(self, label, nodes_labels, value):
        PortHamiltonianObject.__init__(self,label)
        diagRmu, diagQmu = ComputeFractionalDerivativeWeights(value[0], value[1], NbPoles=value[2], NbFreqPoints=200, DoPlot=True,PolesMinMax = (-5,5),out_integrator = True)
        nbPoles = diagRmu.__len__()
        nodes = [nodes_labels[0],]+['node'+label+str(n) for n in range(1,nbPoles)] + [nodes_labels[1],]
        nodes.reverse()
        for n in range(nbPoles):
            N1 = nodes.pop()
            N2 = nodes[-1]

            self.AddLinearDissipativeComponents(['w'+label+str(n)],[diagRmu[n]], [1])
            edge_data_dic = {'ref':'w'+label+str(n), 'type':'dissipative', 'realizability':'effort_controlled', 'linear':True, 'link_ref':'w'+label+str(n)}
            self.Graph.add_edges_from([(N1, N2, edge_data_dic)])

            self.AddLinearStorageComponents(["x"+label+str(n)],[diagQmu[n]], [1])
            edge_data_dic = {'ref':"x"+label+str(n), 'type':'storage', 'realizability':'flux_controlled', 'linear':True, 'link_ref':"x"+label+str(n)}
            self.Graph.add_edges_from([(N1,N2, edge_data_dic)])



def ComputeFractionalDerivativeWeights(p, beta, NbPoles = 20, PolesMinMax = (-5,5),  NbFreqPoints = 2000, FreqMinMax = (10, 30e3), DoPlot = False, out_integrator = False ):

    from scipy.optimize import minimize as OptiMinimize

    if out_integrator:
        beta = beta
    else:
        beta = 1-beta
##############################################################################
    # Func to approximate
    Tbeta = lambda s: s**(-beta)
    
    ##############################################################################
    # Unpack min and max exponents to define the poles 
    emin, emax = PolesMinMax
    Xi  = np.logspace(emin, emax, NbPoles) # xi_0 -> xi_{N+1}
    
    ##############################################################################
    #Frequency grid
    fmin, fmax = FreqMinMax
    wmin, wmax = 2*np.pi*fmin, 2*np.pi*fmax 
    w = np.exp( np.log(wmin) + np.linspace(0,1,NbFreqPoints+1)*np.log(wmax/wmin) )
    w12 = np.sqrt(w[1:]*w[:-1])
    
    ##############################################################################
    # Return the basis vector of elementary damping with poles Xi
    Basis = lambda s, Xi: (s+Xi)**-1
     
    ##############################################################################
    # Prepare Optim
    M = np.zeros((NbFreqPoints,NbPoles), dtype = np.complex64)
    
    for k in np.arange(NbFreqPoints):
        M[k,:] = Basis(1j*w12[k],Xi)
    
    T = Tbeta(1j*w12)
    
    # Perceptual weights
    WBuildingVector = (np.log(w[1:])-np.log(w[:-1]))/(np.abs(T)**2)
    W = np.diagflat(WBuildingVector)
    
    FuncToMinimize = lambda mu: (np.dot(np.conjugate((np.dot(M,mu) - T).T),np.dot(W,np.dot(M,mu) - T))).real
    
    # Constraints
    bnds = [(0,None) for n in range(NbPoles)]
#    cstrs = [{'type': 'ineq', 'fun': lambda mu:  mu[n]} for n in np.arange(NbPoles)]
    
    # Optim
    MuOpt = OptiMinimize(FuncToMinimize, np.ones(NbPoles), bounds = bnds)

    diagQ = []
    diagR = []
    Mu = MuOpt.x
    
    

    if True:
        from matplotlib.pyplot import figure, subplot, loglog, semilogx, ylabel, legend, grid, xlabel     
        TOpt = np.array(M*np.matrix(MuOpt.x).T)
        wmin, wmax = 2*np.pi*20, 2*np.pi*20e3
        figure()
        subplot(2,1,1)
        faxis = w12[(wmin<w12)&(w12<wmax)]/(2*np.pi)
        v1 =20*np.log10( np.abs(T[(wmin<w12)&(w12<wmax)])/beta)
        v2 = 20*np.log10(np.abs(TOpt[(wmin<w12)&(w12<wmax)])/beta)
        v3 = map(lambda x,y: np.abs(x)/np.abs(y),TOpt[(wmin<w12)&(w12<wmax)], T[(wmin<w12)&(w12<wmax)])
        semilogx(faxis,v1,label = 'True') 
        semilogx(faxis,v2, label = 'Approx')
        ylabel('Transfert (dB)')
        legend(loc = 0)
        grid()        
        subplot(2,1,2)
        loglog(faxis, v3, label = 'Error')
        xlabel('Log-frequencies (log Hz)')
        ylabel('Error (dB)')
        legend(loc = 0)
        grid()        
    
    if out_integrator:
        for n in np.arange(NbPoles):
            if Mu[n]>0:
                diagR.append(p*Xi[n]/Mu[n])
                diagQ.append(Mu[n]/p)

    else:
        for n in np.arange(NbPoles):
            if Mu[n]>0:
                diagR.append(Mu[n]/p)
                diagQ.append(Mu[n]*Xi[n]/p)
        
    return diagR, diagQ