//
//  cpldosc_codegen.cpp
//
//  Warren Weckesser
//  Department of Mathematics
//  Colgate University
//  Hamilton, NY
//


#include <iostream>
#include <string>
#include <ginac/ginac.h>


using namespace std;
using namespace GiNaC;

#include "ginac_linalg_funcs.h"

int WriteFunction(string mode, string name, ex expr, lst vars, lst params, lst aux);
int WriteMatrixFunction(string mode, string name, matrix J, lst vars, lst params, lst aux);
int WriteVectorFunction(string mode, string name, matrix v, lst vars, lst params, lst aux);
int WriteGSLVectorField(string name, matrix v, lst vars, lst params, lst aux);
int WriteGSLVectorFieldJacobian(string name, matrix v, lst vars, lst params, lst aux);

///////////////////////////////////////////////////////////
//  main
///////////////////////////////////////////////////////////

int main(int argc, char **argv)
    {
    cout << "--- Code Generator for the Coupled Oscillator Fast/Slow System ---\n";

    // Slow variables
    symbol q1("q1"),q2("q2");

    // Fast variables
    symbol v1("v1"),v2("v2");

    // Parameters (some will be given fixed values later)
    symbol a("a"),sigma1("sigma1"),sigma2("sigma2"),omega("omega");
    symbol r("r"),s("s"),gamma("gamma"),theta("theta");

    lst allparams = lst(a,sigma1,sigma2,omega,r,s,gamma,theta);
    lst fastvars  = lst(v1,v2);
    lst slowvars  = lst(q1,q2);

    // The expression c is the "coupling function"
    symbol v("v");
    ex c = 1/(1+exp(-4*gamma*(v-theta)));

    ex f1 = -(v1-a*tanh(sigma1*v1/a) + q1 + omega*c.subs(v==v2)*(v1-r));
    ex f2 = -(v2-a*tanh(sigma2*v2/a) + q2 + omega*c.subs(v==v1)*(v2-r));
    ex g1 = -q1 + s*v1;
    ex g2 = -q2 + s*v2;

    //
    // Fix many of the parameter values here.  Only sigma1, sigma2 and omega
    // will remain as parameters.
    //
    lst values = lst(a==1,s==1,gamma==10,r==-4,theta==numeric(1)/75);
    matrix f = matrix(2,1,lst(f1.subs(values),f2.subs(values)));
    matrix g = matrix(2,1,lst(g1.subs(values),g2.subs(values)));
    lst params = lst(sigma1,sigma2,omega);

    matrix dfv = jacobian(f,fastvars);

    symbol f11("f11"), f12("f12"), f21("f21"), f22("f22");

    cout << "--- Writing cpldosc_foldfunc.c (C)\n";

    WriteFunction("c","cpldosc_foldfunc",f11*f22-f12*f21,fastvars,params,
              lst(f11==dfv(0,0), f12==dfv(0,1), f21==dfv(1,0), f22==dfv(1,1) )); 

    symbol df11v1("df11v1"),df11v2("df11v2");
    symbol df12v1("df12v1"),df12v2("df12v2");
    symbol df21v1("df21v1"),df21v2("df21v2");
    symbol df22v1("df22v1"),df22v2("df22v2");

    matrix detjac(2,1);
    detjac(0,0) = f11*df22v1 + df11v1*f22 - f12*df21v1 - df12v1*f21;
    detjac(1,0) = f11*df22v2 + df11v2*f22 - f12*df21v2 - df12v2*f21;

    cout << "--- Writing cpldosc_foldfunc_grad.c (C)\n";

    WriteVectorFunction("c","cpldosc_foldfunc_grad",detjac,fastvars,params,
              lst(f11==dfv(0,0), f12==dfv(0,1), f21==dfv(1,0), f22==dfv(1,1),
                  df11v1==dfv(0,0).diff(v1),df11v2==dfv(0,0).diff(v2),
                  df12v1==dfv(0,1).diff(v1),df12v2==dfv(0,1).diff(v2),
                  df21v1==dfv(1,0).diff(v1),df21v2==dfv(1,0).diff(v2),
                  df22v1==dfv(1,1).diff(v1),df22v2==dfv(1,1).diff(v2)));

    //
    // Solve for the slow variables.
    //
    lst eqns;
    eqns = f(0,0)==0, f(1,0)==0;
    ex sol = lsolve(eqns,slowvars);

    matrix ghat(2,1);
    ghat(0,0) = g(0,0).subs(sol);
    ghat(1,0) = g(1,0).subs(sol);

    cout << "--- Writing cpldosc_ghat.c (C)\n";

    WriteVectorFunction("c","cpldosc_ghat",ghat,fastvars,params,lst());

    matrix fhat(2,1);
    fhat(0,0) = sol[0].rhs();
    fhat(1,0) = sol[1].rhs();

    cout << "--- Writing cpldosc_fhat.c (C)\n";

    WriteVectorFunction("c","cpldosc_fhat",fhat,fastvars,params,lst());

    matrix dfhat = jacobian(fhat,fastvars);

    cout << "--- Writing cpldosc_dfhat.c (C)\n";

    WriteMatrixFunction("c","cpldosc_dfhat",dfhat,fastvars,params,lst());

    matrix adj_dfhat = adjoint(dfhat);
    matrix slow_vf(2,1);
    symbol adj11("adj11"), adj12("adj12"), adj21("adj21"), adj22("adj22");
    symbol ghat1("ghat1"), ghat2("ghat2");
    slow_vf(0,0) = adj11 * ghat1 + adj12 * ghat2;
    slow_vf(1,0) = adj21 * ghat1 + adj22 * ghat2;

    cout << "--- Writing cpldosc_slowvf.c, a GSL Vector Field\n";

    WriteGSLVectorField("cpldosc_slowvf",slow_vf,fastvars,params,
         lst(adj11==adj_dfhat(0,0),adj12==adj_dfhat(0,1),
             adj21==adj_dfhat(1,0),adj22==adj_dfhat(1,1),
             ghat1==ghat(0,0), ghat2==ghat(1,0)) );

    matrix slow_vf_jac(2,2);
    slow_vf_jac(0,0) = adj11*ghat(0,0).diff(v1) + adj_dfhat(0,0).diff(v1)*ghat1 +
                       adj12*ghat(1,0).diff(v1) + adj_dfhat(0,1).diff(v1)*ghat2;
    slow_vf_jac(0,1) = adj11*ghat(0,0).diff(v2) + adj_dfhat(0,0).diff(v2)*ghat1 +
                       adj12*ghat(1,0).diff(v2) + adj_dfhat(0,1).diff(v2)*ghat2;
    slow_vf_jac(1,0) = adj21*ghat(0,0).diff(v1) + adj_dfhat(1,0).diff(v1)*ghat1 +
                       adj22*ghat(1,0).diff(v1) + adj_dfhat(1,1).diff(v1)*ghat2;
    slow_vf_jac(1,1) = adj21*ghat(0,0).diff(v2) + adj_dfhat(1,0).diff(v2)*ghat1 +
                       adj22*ghat(1,0).diff(v2) + adj_dfhat(1,1).diff(v2)*ghat2;

    cout << "--- Writing cpldosc_slowvf_jac.c, a GSL Vector Field Jacobian\n";

    WriteGSLVectorFieldJacobian("cpldosc_slowvf_jac",slow_vf_jac,fastvars,params,
         lst(adj11==adj_dfhat(0,0),adj12==adj_dfhat(0,1),
             adj21==adj_dfhat(1,0),adj22==adj_dfhat(1,1),
             ghat1==ghat(0,0), ghat2==ghat(1,0)) );

    }