#include "IBeamlineCoordinates.hxx"
#include "IBeamlineParameterExtractor.hxx"
#include <ICOMETLog.hxx>
#include <HEPUnits.hxx>
#include <TVector3.h>
#include <TLine.h>
#include <TFile.h>
#include <TH1.h>
#include <TH2.h>
#include <TGraph.h>
#include <TGeoManager.h>
#include <TGeoMaterial.h>
#include <TGeoNode.h>
#include <TGeoVolume.h>
#include <TMath.h>
#include <cstdlib>
#include <sstream>
#include <IFieldMapList.hxx>
#include <IFieldManager.hxx>
#include <IFieldMapLoader.hxx>
using std::cout;
using std::endl;

double GetDensity(const double* point);

std::ostream& operator<<(std::ostream& stream,const TVector3& vect){
    stream<<"("<<vect.x()<<", "<<vect.y()<<", "<<vect.z() <<")";
    return stream;
}

void Usage(const char* prog, std::string inputMessage){
    cout<<"Usage: "<<prog<<" [options]\n"
          "\n" << inputMessage << "\n"
          "Output file Options: \n"
          " -n <n_beamline_points>    Set the number of points along the beam axis that should be checked.\n"
          " -H <height>      Set the maximum height above and below the beamline to sample.\n"
          " -d <dist_radial> Set the transverse radial distance from the beam axis to scan along.\n"
          " -p <phi>         Set the transverse azimuthal angle to the vertical to scan along (in radians).\n"
          " -o <outfile>     Set the output file name.\n"
          "\n"
          "Beamline Options: \n"
          " -b <beamline>    Set the beamline type.\n"
          " -B <filename>    Set the beamline to match the geometry in the named file.\n"
          "Other Options: \n"
          " -I <filename>    Load geometry and fieldmap from a root file. Equivalent to: -B <filename> -R <filename>\n"
        <<endl;
}

int main(int argc,char** argv){
    COMET::ICOMETLog::Configure();

    double n_beamline_points=500, transverseAngle=0,  transverseDist=0;
    std::string outfilename="beamline_field.root";
    COMET::IFieldMapList* field_list=new COMET::IFieldMapList;
    COMET::IFieldMapLoader loader;
    COMET::IBeamlineCoordinates beamline;
    double top=600;
    for (;;) {
        std::string optString = "h?"+loader.GetOptString()+"n:d:p:o:b:B:I:H:";
        char c = getopt(argc, argv, optString.c_str());
        std::stringstream stream;
        if(optarg) stream.str(optarg);
        // Try to process it in the loader
        int inputArg = loader.ProcessArgument(c, optarg, field_list);
        // Process R option again
        if(inputArg == 0 && c == 'R') inputArg = 3;
        if(inputArg == 0) return 0;
        // Bad argument, pass to main
        else if (inputArg == 1 || inputArg == 2) return inputArg;
        else if (inputArg == 3){
            if (c<0) break;
            switch (c) {
                // output options
                case 'n':
                    {
                        stream>>n_beamline_points;
                        break;
                    }
                case 'H':
                    {
                        stream>>top;
                        break;
                    }
                case 'p':
                    {
                        stream>>transverseAngle;
                        break;
                    }
                case 'd':
                    {
                        stream>>transverseDist;
                        break;
                    }
                case 'o':
                    {
                        stream>>outfilename;
                        break;
                    }
                case 'b':
                    {
                        std::string beam_type(optarg);
                        if(beam_type=="g4bl_p2"){
                            COMETLog("Setting beamline to g4beamline phase-II parameters, (ie. default)");
                        }else if (beam_type=="111216"){
                            COMETLog("Setting beamline to phase-II parameters matching fields in 111216");
                            beamline.SetLength_prodTgtSec(4875);
                            beamline.SetLength_Tor1_exit(800);
                        }else if (beam_type=="150630_p1"){
                            COMETLog("Setting beamline to phase-I parameters matching fields in 150630");
                            beamline.SetPhase(1);
                            beamline.SetLength_Tor1_exit(700);
                            beamline.SetLength_prodTgtSec(5175);
                        }else{
                            COMETError("Bad beamline specification: "<<beam_type);
                            COMETError("Should be either: 'g4bl_p2', '111216' "<<beam_type);
                            return 1;
                        }
                        continue;
                        break;
                    }
                case 'B':
                    {
                        COMET::IBeamlineParameterExtractor::Load(beamline,TGeoManager::Import(optarg));
                        continue;
                    }
                case 'I':
                    {
                        if(!COMET::IBeamlineParameterExtractor::Load(beamline,TGeoManager::Import(optarg))
                        || !COMET::IFieldManager::Import(optarg)){
                            return 1;
                        }
                        continue;
                    }
                case 'h':case '?':
                    Usage(argv[0], loader.UsageMessage());
                    return 0;
                default:
                    cout<<"Error: Unknown option: "<<c<<endl<<endl;
                    Usage(argv[0], loader.UsageMessage());
                    return 1;
            }
        }
        if(inputArg == 4 || stream.fail() || ! stream.eof()){
            cout<<"Error: Bad option value: '"<<optarg<<"' for option '"<<c <<"'"<<endl<<endl;
            Usage(argv[0], loader.UsageMessage());
            return 1;
        }
    }
    // Need to init the global field object
    COMET::IFieldManager* mgr=COMET::IFieldManager::GetObject();
    field_list->ls();
    field_list->InitialiseField();
    beamline.Initialise();


    TFile* outfile=TFile::Open(outfilename.c_str(),"recreate");

    const double length=beamline.GetTotalLength();
    const double lengthPerbin=length/n_beamline_points;
    int binsUp=top/lengthPerbin;
    top=binsUp*lengthPerbin;
    COMETError(binsUp<<", "<<top);
    TH1* hMagnitude      = new TH1F("hMagnitude"      , "Magnitude on field on axis"             ,n_beamline_points+1,0,length);
    TH1* hVertical       = new TH1F("hVertical"       , "Vertical component of field on axis"    ,n_beamline_points+1,0,length);
    TH1* hHorizontal     = new TH1F("hHorizontal"     , "Horizontal component of field on axis"  ,n_beamline_points+1,0,length);
    TH1* hLongitudinal   = new TH1F("hLongitudinal"   , "Longitudinal component of field on axis",n_beamline_points+1,0,length);
    TH1* hMarkers        = new TH1F("hMarkers"        , "Longitudinal component of field on axis",n_beamline_points+1,0,length);
    TGraph* outline=new TGraph();

    TH2* hMagnitude2d    = new TH2F("hMagnitude2d"    , "Magnitude on field on axis"             ,n_beamline_points+1,0,length,2*binsUp+1,-top,top);
    TH2* hVertical2d     = new TH2F("hVertical2d"     , "Vertical component of field on axis"    ,n_beamline_points+1,0,length,2*binsUp+1,-top,top);
    TH2* hHorizontal2d   = new TH2F("hHorizontal2d"   , "Horizontal component of field on axis"  ,n_beamline_points+1,0,length,2*binsUp+1,-top,top);
    TH2* hLongitudinal2d = new TH2F("hLongitudinal2d" , "Longitudinal component of field on axis",n_beamline_points+1,0,length,2*binsUp+1,-top,top);
    TH2* hGeometry       = new TH2F("hGeometry"       , "Geometry along the beamline axis"       ,n_beamline_points+1,0,length,2*binsUp+1,-top,top);

    const double beamline_step=beamline.GetTotalLength()/n_beamline_points;
    double point=0;
    double global_arr[3];
    double field_vectors[6];
    TVector3 vertical(0,1,0);
    for(int i=0;i<n_beamline_points; ++i){
        // coords for this point
        const TVector3 global=beamline.ToGlobalCoords(point,transverseDist,transverseAngle);
        const TVector3 direction=beamline.GetBeamlineDirectionFromBeamline(point);
        const TVector3 horizontal=direction.Cross(vertical);

        for(int j=-binsUp; j<binsUp; ++j){
            double height=j*lengthPerbin;

            // Get field
            global_arr[0]=global.X(); global_arr[1]=global.Y()+height; global_arr[2]=global.Z(); 
            mgr->GetFieldValue(global_arr,field_vectors);
            TVector3 field(field_vectors[0],field_vectors[1],field_vectors[2]);
            field*=1/unit::tesla;

            if(j==0){
                //COMETError("Filling j==0 plots, i=="<<i);
                hMagnitude->Fill(point,field.Mag());
                hLongitudinal->Fill(point,field.Dot(direction));
                hVertical->Fill(point,field.Dot(vertical));
                hHorizontal->Fill(point,field.Dot(horizontal));
                outline->SetPoint(i,-global.X(),global.Z());
            }
            hMagnitude2d->Fill(    point, height, field.Mag()           ); 
            hVertical2d ->Fill(    point, height, field.Dot(vertical)   ); 
            hLongitudinal2d->Fill( point, height, field.Dot(direction)  ); 
            hHorizontal2d->Fill(   point, height, field.Dot(horizontal) ); 
            if(gGeoManager){
                const double density=GetDensity(global_arr);
                hGeometry->Fill(point,height,density);
            }
        }

        point+=beamline_step;
    }

    std::vector<double> component_edges;
    beamline.GetComponentEnds(component_edges);
    double max=hMagnitude->GetMaximum();
    for(auto i_comp=component_edges.begin(); i_comp!=component_edges.end(); ++i_comp){
        TVector3 left=beamline.ToGlobalCoords(*i_comp,-100,TMath::Pi()/2);
        TVector3 right=beamline.ToGlobalCoords(*i_comp,100,TMath::Pi()/2);
        TLine* line=new TLine(-left.X(),left.Z(),-right.X(),right.Z());
        outline->GetListOfFunctions()->Add(line);
        line=new TLine(*i_comp,0,*i_comp,max/2);
        hMarkers->GetListOfFunctions()->Add(line);
    }
    outline->SetLineWidth(2);
    outline->SetLineColor(kBlack);

    outline->Write("outline");
    outfile->Write();
    outfile->Close();

    return 0;
}

double GetDensity(const double* point){
        // Factors of 10 come from units in geometry being cm, not mm as used by
        // field
        const TGeoNode* node=gGeoManager->FindNode(point[0]/10,point[1]/10,point[2]/10);
        if(!node){
            COMETDebug("No node, perhaps in the world volume?");
            return -1;
        }
        const TGeoVolume* volume =node->GetVolume();
        if(!volume){
            COMETError("Node without a volume !");
            return -1;
        }
        const TGeoMaterial* material =volume->GetMaterial();
        return material->GetDensity();
}