#include <vector>
#include <iostream>

#include <TH1.h>
#include <TMarker.h>
#include <TArrow.h>
#include <TCanvas.h>
#include <TRandom.h>

#include "MFresnelLens.h"

/*****************************************************************

 fresnellens_traceray.C - 2D Example of ray-tracing the fresnel lens

 The macro has to be compilesd to work

 To run the macro from the command line (assuming you are in a directory
 Mars/build where you have built your Mars environment) you have to do

    root ../hawc/fresnellens_traceray.C++

 or from within root

    [0] .x ../hawc/fresnellens_traceray.C++

******************************************************************/
void fresnellens_traceray()
{
    // ------------------- setup lens -----------------------

    const double D = 54.92;  // [cm] Diameter of the refractive surface
    const double F = 50.21;  // [cm] Nominal focal length of the lens

    const double H = 0.25;   // [cm] Width of a single groove
    const double w = 0.01;   // [cm] Width of a single groove

    // For a well visible structure, use instead
    //const double H = 2.5;   // [cm] Width of a single groove
    //const double w = 1.0;   // [cm] Width of a single groove

    //const double w = 1;        // [cm] Width of a single groove
    //const double H = 2.5;    // [cm] Thickness of lens

    const double Z = F+H;    // [cm] camera position w.r.t. lens exit (see also MGeomCamFAMOUS!)
    const double R = D/2;    // [cm] radius of lens

    const double lambda0 = 546; // [nm] Wavelength for lens surface definition
    const double lambda  = 546; // [nm] Wavelength for the simulated photons

    const bool point_source = false;  // Enable simulation of a point source

    const double angle = 6; // [deg] Angle of incidence of the simulated rays

    const double Z0 = point_source ? F : 3; // [cm] Starting Z-position of rays
    const double X0 = Z0*atan(angle*TMath::DegToRad()); // [cm] (If<0: parallel rays (see angle), If>=0: point source at X0/Z0)

    MFresnelLens lens;
    lens.DefineLens(F, D, w, H, lambda0);
    //lens.SetPSF(psf);
    //lens.EnableSlopeAbsorption();
    //lens.EnableDraftAbsorption();
    //lens.DisableBottomReflection();
    //lens.DisableMultiEntry();
    //lens.DisableFresnelReflection();
    lens.ReadTransmission("resmc/hawcseye/transmission-pmma-3mm.txt", 0.3, true);

    // ------------------ Setup display ---------------------

    TH1C h("", "", 10000, -R*1.1, R*1.1);
    h.SetStats(kFALSE);
    h.SetMaximum( 3);
    h.SetMinimum(-1.1*F);
    h.DrawCopy();

    TLine line;
    line.SetLineWidth(2);

    for (unsigned int i=0; i<lens.GetNumGrooves(); i++)
    {
        const MFresnelLens::Groove &g = lens[i];

        line.SetLineColor(kBlack);
        line.DrawLine( w*i, 0,  g.r, g.slope.h);
        line.DrawLine(-w*i, 0, -g.r, g.slope.h);;

        line.SetLineColor(kRed);
        line.DrawLine( g.r, g.slope.h,  w*i+w, 0);
        line.DrawLine(-g.r, g.slope.h, -w*i-w, 0);
    }

    line.SetLineColor(kBlack);
    line.SetLineWidth(1);
    line.DrawLine(-30, -H,   30, -H);
    line.DrawLine(-30, -Z, 30, -Z);
    line.DrawLine( R,  -Z,  R,  0);
    line.DrawLine(-R,  -Z, -R,  0);
    line.SetLineWidth(3);
    line.DrawLine(-4.5*1.5, -Z, 4.5*1.5, -Z);
    line.SetLineWidth(1);

    // ------------------------------------------------------

    if (point_source)
        cout << "\nPoint source at x=" << X0 << " z=" << Z0 << "\n" << endl;

    // ------------------------------------------------------

    TArrow arrow;
    TMarker marker;

    marker.SetMarkerColor(kBlue);
    marker.SetMarkerStyle(kFullDotMedium);

    arrow.SetArrowSize(0.01);
    arrow.SetLineColor(kBlue);

    double vc = 1./(TMath::C()*100/1e9); // cm/ns

    for (int i=0; i<100; i++)
    {
        double X = gRandom->Uniform(-R, R);

        TVector3 pos(X, 0,  0);
        TVector3 dir(0, 0, -1);

        double theta = point_source ?
            TMath::Pi()-atan((X+X0)/Z0) :
            (180-angle)*TMath::DegToRad();

        dir.SetMagThetaPhi(1, theta, 0);

        MQuaternion p(pos, 0);
        MQuaternion u(dir, vc);

        // Propagate to starting plane in front of the lens
        p.PropagateZ(u, Z0);

        vector<MQuaternion> vec;
        vec.push_back(p);

        const int cnt = lens.TraceRay(vec, p, u, lambda, true);

        // Particle successfully reached the focal plane
        if (cnt>=0)
        {
            p.PropagateZ(u, -Z);
            vec.push_back(p);

            p.PropagateZ(u, u.Z()<0 ? -2.5*F : Z0);
            vec.push_back(p);
        }

        // Particle is upgoing and has hit no surface
        if (cnt==-1291 || cnt==-1391)
        {
            p.PropagateZ(u, 3);
            vec.push_back(p);
        }

        arrow.SetLineColor(kBlue);
        for (unsigned int j=0; j<vec.size()-1; j++)
        {
            arrow.DrawArrow(vec[j].X(), vec[j].Z(), vec[j+1].X(), vec[j+1].Z());
            marker.DrawMarker(vec[j+1].X(), vec[j+1].Z());
        }
    }
}