#include <iostream>
#include <fstream>
#include <cmath>

#include "rootracker_info.h"

using std::cout;
using std::endl;

///////////////////////////////////////////////////////////////////////////////
//                                                                           //
//  Coordinate Systems (angles are not to scale)                             //
// SPATIAL:                                                                  //
//                                     __ horizontal                         //
//                                 __--                                      //
//                  vertical . __--_______ 10 degrees                        //
//             \               \                                             //
//    proton    \               \ longitudinal                               //
//     beam      \               \                                           //
//                \                                                          //
//                 \                                                         //
//                 ---                                                       //
//     10 deg   \  \  \                                                      //
//            ---\  \  \ target     z  <------| . y (out of screen)          //
//                   \__\                     |                              //
//                                            |                              //
//                                            |                              //
//                                            v x                            //
//                                                                           //
//   Coordinate Transformations:                                             //
//       theta = 170 degrees                                                 //
//           y = vertical                                                    //
//        (-z) = (cos(theta) -sin(theta)) (h)                                //
//         (x) = (sin(theta)  sin(theta)) (l)                                //
//     => z = -(horizontal*cos(theta) - longitudinal*sin(theta))             //
//     => x = horizontal*sin(theta) + longitudinal*cos(theta)                //
//                                                                           //
// ANGULAR:                                                                  //
//                  vertical                                                 //
//                     ^                                                     //
//                     |                                                     //
//                     |   momentum                                          //
//                     |  /                                                  //
//                     | /|                                                  //
//                     |/_|________> horizontal                              //
//                    /\theta vert                                           //
//                   /  \ |                                                  //
//  longitudunal    / ^  \|                                                  //
//                 / theta                                                   //
//                   horiz                                                   //
//        p_vert = mom*sin(theta_vert)                                       //
//        p_horiz = mom*cos(theta_vert)*sin(theta_horiz)                     //
//        p_longitudinal = mom*cos(theta_vert)*cos(theta_horiz)              //
//                                                                           //
///////////////////////////////////////////////////////////////////////////////

// EDIT (AE): 15th July 2014: looking at the result, it seems x and z need to be swapped. I'm not entirely sure how to show this in the diagrams above but thought I should mention it.

void ConvertHVLToXYZ(const std::vector<double>& input, std::vector<double>& output);

void Usage(const char* program);

int main(int argc, char* argv[]) {
    if(argc<=1 || std::string(argv[1])=="-h"){
        Usage(argv[0]);
        return 0;
    }

    TFile* file = NULL;
    TTree* tree = NULL;
    bool set_up = false;

    std::string full_input_path = argv[1];

    // replace extension with .root to give output name
    std::string full_output_path = full_input_path;
    size_t dot=full_input_path.rfind('.');
    if(dot!=std::string::npos){
        full_output_path.replace(dot+1, std::string::npos, "root"); 
    } else {
        full_output_path+=".root";
    }

    // Open the input text file
    std::ifstream input(full_input_path.c_str(), std::ifstream::in);
    if (!input.is_open()) {
        std::cout << "Problem opening " << full_input_path << std::endl;
        return 1;
    }

    // Set up all the branches
    // this should probably be moved to a separate function in oaRooTracker
    if (!set_up) {

        file = new TFile(full_output_path.c_str(), "RECREATE");
        tree = new TTree("gRooTracker", "Proton Beam event tree rootracker format");

        tree->Branch("EvtFlags", "TBits",      &brEvtFlags, 32000, 1);           
        tree->Branch("EvtCode",  "TObjString", &brEvtCode,  32000, 1);           
        tree->Branch("EvtNum",          &brEvtNum,          "EvtNum/I");             
        tree->Branch("EvtXSec",         &brEvtXSec,         "EvtXSec/D");           
        tree->Branch("EvtDXSec",        &brEvtDXSec,        "EvtDXSec/D");           
        tree->Branch("EvtWght",         &brEvtWght,         "EvtWght/D");           
        tree->Branch("EvtProb",         &brEvtProb,         "EvtProb/D");           
        tree->Branch("EvtVtx",           brEvtVtx,          "EvtVtx[4]/D");             
        tree->Branch("StdHepN",         &brStdHepN,         "StdHepN/I");             
        tree->Branch("StdHepPdg",        brStdHepPdg,       "StdHepPdg[StdHepN]/I"); 
        tree->Branch("StdHepStatus",     brStdHepStatus,    "StdHepStatus[StdHepN]/I");
        tree->Branch("StdHepRescat",     brStdHepRescat,    "StdHepRescat[StdHepN]/I");
        tree->Branch("StdHepX4",         brStdHepX4,        "StdHepX4[StdHepN][4]/D");
        tree->Branch("StdHepP4",         brStdHepP4,        "StdHepP4[StdHepN][4]/D");
        tree->Branch("StdHepPolz",       brStdHepPolz,      "StdHepPolz[StdHepN][3]/D");
        tree->Branch("StdHepFd",         brStdHepFd,        "StdHepFd[StdHepN]/I");
        tree->Branch("StdHepLd",         brStdHepLd,        "StdHepLd[StdHepN]/I");
        tree->Branch("StdHepFm",         brStdHepFm,        "StdHepFm[StdHepN]/I");
        tree->Branch("StdHepLm",         brStdHepLm,        "StdHepLm[StdHepN]/I");

        // Added for COMET
        tree->Branch("Weight",         brWeight,        "Weight[StdHepN]/D");
        tree->Branch("MonitorName", "TObjArray", brMonitorName);

        set_up = true;
    }

    // Get the variables ready for reading in Note that this coordinate system
    // is looking down the proton beam line to the target with the origin 8500mm
    // along the proton beam line away from the target I have defined this
    // coordinate system as (horizonal, vertical, longitudinal)
    // coords [cm], angles [mrad], momentum [GeV]
    double horizontal_coord, horizontal_angle, vertical_coord, vertical_angle, momentum, weight, event_number;
    double longitudinal_origin = 850.0; // cm

    while (input.good()) {

        input >> horizontal_coord >> horizontal_angle >> vertical_coord >> vertical_angle >> momentum >> weight >> event_number;

        if (input.eof()) { // check if we have the EOF, otherwise we get the last line twice
            break;
        }

        // Convert angles to rads from mrads
        horizontal_angle /= 1000;
        vertical_angle /= 1000;

        // Convert position from HVL to XYZ
        std::vector<double> hvl, xyz;
        hvl.push_back(horizontal_coord);
        hvl.push_back(vertical_coord);
        hvl.push_back(longitudinal_origin);
        ConvertHVLToXYZ(hvl, xyz);

        double x = -(xyz.at(2));
        double y = xyz.at(1);
        double z = -(xyz.at(0));

        double momentum_vertical = momentum*std::sin(vertical_angle);
        double momentum_horizontal = momentum*std::cos(vertical_angle)*std::sin(horizontal_angle);
        double momentum_longitudinal = momentum*std::cos(vertical_angle)*std::cos(horizontal_angle);

        // Convert the momentum from HVL to XYZ
        hvl.clear(); xyz.clear();
        hvl.push_back(momentum_horizontal);
        hvl.push_back(momentum_vertical);
        hvl.push_back(momentum_longitudinal);
        ConvertHVLToXYZ(hvl, xyz);

        double p_x = xyz.at(2); // want the momentum to point in the opposite direction
        double p_y = xyz.at(1);
        double p_z = xyz.at(0);

        double mass = 0.9382720; // proton mass in GeV
        double E = std::sqrt(momentum*momentum + mass*mass);

        /*
           std::cout << horizontal_coord << " " << horizontal_angle << " " << vertical_coord << " " << vertical_angle << " " << momentum << " " << weight << " " << event_number << std::endl;
           std::cout << "(x, y, z) = (" << x << ", " << y << ", " << z << ")" << std::endl;
           std::cout << "(p_h, p_v, p_l) = (" << momentum_horizontal << ", " << momentum_vertical << ", " << momentum_longitudinal << ")" << std::endl;
           std::cout << "(p_x, p_y, p_z) = (" << p_x << ", " << p_y << ", " << p_z << ")" << std::endl;
           std::cout << "E = " << E << std::endl;
           */

        // Now add to tree
        brEvtNum = event_number;
        brEvtXSec = -999;
        brEvtDXSec = -999;
        brEvtWght = -999;
        brEvtProb = -999;
        brEvtVtx[0] = 0; brEvtVtx[1] = 0; brEvtVtx[2] = 0; brEvtVtx[3] = 0; // initialise to zero since I think this is the origin of the event
        brStdHepN = 1;

        brStdHepPdg[0] = 2212; // a proton
        brStdHepStatus[0] = 1; // set to 1 since that is what SimG4 is expecting for particles that it can track
        brStdHepRescat[0] = -999;  
        brStdHepX4    [0][0] = x; // don't convert from cm
        brStdHepX4    [0][1] = y;
        brStdHepX4    [0][2] = z;
        brStdHepX4    [0][3] = 0; // set time to 0?

        brStdHepP4    [0][0] = p_x; // already in GeV for SimG4
        brStdHepP4    [0][1] = p_y;
        brStdHepP4    [0][2] = p_z;
        brStdHepP4    [0][3] = E;

        brStdHepPolz  [0][0] = -999;
        brStdHepPolz  [0][1] = -999;
        brStdHepPolz  [0][2] = -999;

        brStdHepFd    [0] = -999;     // First daughter
        brStdHepLd    [0] = -999;     // Last  daughter
        brStdHepFm    [0] = -999;     // First mother
        brStdHepLm    [0] = -999;     // Last  mother

        brWeight      [0] = weight;

        tree->Fill();
    }

    file->Write();
    file->Close();
    return 0;
}

/// ConvertHVLToXYZ
/// Converts (horizontal, vertical, longitudinal) to (x, y,z)
void ConvertHVLToXYZ(const std::vector<double>& input, std::vector<double>& output) {

    double target_rotation = 10; // target is rotated 10 degrees
    double pi = 3.14159265359;
    double coordinate_rotation = (180 - target_rotation)*(pi/180); // want to rotate 170 degrees anti-clockwise


    double h = input.at(0);
    double v = input.at(1);
    double l = input.at(2);

    double x = h*std::sin(coordinate_rotation) + l*std::cos(coordinate_rotation);
    double y = v;
    double z = -(h*std::cos(coordinate_rotation) - l*std::sin(coordinate_rotation));

    output.push_back(x);
    output.push_back(y);
    output.push_back(z);
}

void Usage(const char* program){
    cout<<"Usage: "<<program<<" [-h|INPUT_FILE]\n";
    cout<<"\n";
    cout<<"Convert a turtle simulation output file into a rootracker file.\n";
    cout<<"'INPUT_FILE' should specify the full path and name of the input file to convert.\n";
    cout<<"The resultant output file is written to the same directory as the input file.\n";
    cout<<"'-h' will print this help message.\n";
}