#include "IFieldPlotter.hxx"
#include "IFieldManager.hxx"
#include <ICOMETLog.hxx>
#include <sstream>
#include <cmath>
#include <TH2.h>

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

COMET::IFieldPlotter::IFieldPlotter():
    fUpVector(0,0,1),
    fNormal(0,1,0),
    fOrigin(0,0,0),
    fStepUp(1),
    fStepSide(1),
    fNSteps(500),
    fDirectory(TDirectory::CurrentDirectory()) {
    }

bool COMET::IFieldPlotter::Initialise(){
    COMETNamedLog("IFieldPlotter","Setting up with "
            "Normal: " <<fNormal<<
            " Up vector: "<< fUpVector
            );

    fDirOut=fNormal.Unit();
    fDirUp=fUpVector-fUpVector.Dot(fNormal)*fNormal;
    if(fDirUp.Mag()<1e-14) {
        COMETError("Up vector and plane normals should not be the same");
        return false;
    }
    fDirUp=fDirUp.Unit();
    fDirSide=fDirUp.Cross(fDirOut);

    COMETNamedLog("IFieldPlotter","Plane defined by: "
            "Normal: " <<fDirOut<<
            " Up vector: "<< fDirUp<<
            " Side vector:"<<fDirSide
            );
    return true;
}

bool COMET::IFieldPlotter::PlotPlane(TDirectory* output){
    TDirectory* oldwd=TDirectory::CurrentDirectory();
    if(output) output->cd();
    else if(fDirectory) fDirectory->cd();

    COMET::IFieldManager* theGlobalField = COMET::IFieldManager::getObject();

    double limit_side=fNSteps*fStepSide;
    double limit_up=fNSteps*fStepUp;
    std::stringstream x_axis; x_axis<<"Distance along "<<fDirSide;
    std::stringstream y_axis; y_axis<<"Distance along "<<fDirUp;

    std::stringstream title;
    title<<"Magnitude of field through "<<fOrigin 
        <<";"<<x_axis.str()<<";"<<y_axis.str();
    TH2* hMagnitude=new TH2F("hMagnitude",title.str().c_str(),
            2*fNSteps+1 , -limit_side , limit_side , 2*fNSteps+1 , -limit_up , limit_up);

    title.str(""); 
    title<<"Component of field parallel to "<<fDirSide<<" through "<<fOrigin
        <<";"<<x_axis.str()<<";"<<y_axis.str();
    TH2* hSideways=new TH2F("hSideways",title.str().c_str(),
            2*fNSteps+1 , -limit_side , limit_side , 2*fNSteps+1 , -limit_up , limit_up);

    title.str(""); 
    title<<"Component of field parallel to "<<fDirUp<<" through "<<fOrigin
        <<";"<<x_axis.str()<<";"<<y_axis.str();
    TH2* hUpwards=new TH2F("hUpwards",title.str().c_str(),
            2*fNSteps+1 , -limit_side , limit_side , 2*fNSteps+1 , -limit_up , limit_up);

    title.str(""); 
    title<<"Component of field parallel to "<<fDirOut<<" through "<<fOrigin
        <<";"<<x_axis.str()<<";"<<y_axis.str();
    TH2* hOutwards=new TH2F("hOutwards",title.str().c_str(),
            2*fNSteps+1 , -limit_side , limit_side , 2*fNSteps+1 , -limit_up , limit_up);

    double point[3];
    double field_vectors[6];
    for(int i_side = -fNSteps; i_side<=fNSteps; ++i_side){
        double dist_side=fStepSide*i_side;
        TVector3 side= fDirSide*dist_side;
        for(int i_up = -fNSteps; i_up<=fNSteps; ++i_up){
            double dist_up=fStepUp*i_up;
            TVector3 position= fDirUp*dist_up;
            position+= side;
            position+= fOrigin;

            point[0]=position.x();
            point[1]=position.y();
            point[2]=position.z();
            for(int j=0; j<6;++j) field_vectors[j]=0;

            theGlobalField->GetFieldValue(point,field_vectors);
            TVector3 field(field_vectors[0],field_vectors[1],field_vectors[2]);

            hMagnitude->Fill(dist_side,dist_up,field.Mag());
            hUpwards->Fill(dist_side,dist_up,field.Dot(fDirUp));
            hSideways->Fill(dist_side,dist_up,field.Dot(fDirSide));
            hOutwards->Fill(dist_side,dist_up,field.Dot(fDirOut));
        }
    }
    oldwd->cd();
    return true;
}

bool COMET::IFieldPlotter::PlotPlanes(int n_planes,double plane_sep){
    double start_depth=(n_planes-1)*0.5*plane_sep;
    int start=(n_planes-1)*0.5;
    int stop=start + (n_planes+1)%2;
    TVector3 orig_origin=fOrigin;
    fOrigin-=start_depth*fDirOut;

    for(int i_plane=-start;i_plane<=stop; ++i_plane){
        // create a name for the plane
        std::string slice_name=Form("slice_%g_%g_%g",fOrigin.x(),fOrigin.y(),fOrigin.z());
        TDirectory* dir=fDirectory->mkdir(slice_name.c_str());

        COMETNamedLog("IFieldPlotter","Running on plane ("<<i_plane<<") centred at "<<fOrigin);

        PlotPlane(dir);
        fOrigin+=fDirOut*plane_sep;
    }
    fOrigin=orig_origin;
    return true;
}