#include <RAT/GeoSc48SourceFactory.hh>

#include <RAT/DB.hh>
#include <RAT/Log.hh>
#include <RAT/Detector.hh>
#include <RAT/string_utilities.hpp>

#include <G4Material.hh>

#include <G4LogicalVolume.hh>
#include <G4ThreeVector.hh>
#include <G4RotationMatrix.hh>

#include <G4Tubs.hh>
#include <G4Polycone.hh>

#include <G4SubtractionSolid.hh>
#include <G4UnionSolid.hh>

#include <G4PVPlacement.hh>

#include <G4Colour.hh>
#include <G4VisAttributes.hh>

#include <vector>
#include <algorithm>
#include <functional>
#include <CLHEP/Units/SystemOfUnits.h>
#include <CLHEP/Units/PhysicalConstants.h>
#include <string>

using namespace CLHEP;

namespace RAT
{

void GeoSc48SourceFactory::SetColor(DBLinkPtr table, G4String colorName, G4LogicalVolume *logicalVolume)
{
    G4VisAttributes *vis = new G4VisAttributes();
    try {
        const std::vector<double> &color = table->GetDArray(colorName);
        Log::Assert(color.size() == 3 || color.size() == 4,
                    "GeoSc48SourceFactory: Color entry " + colorName + " does not have 3 (RGB) or 4 (RGBA) components");
        if (color.size() == 3) // RGB
            vis->SetColour(G4Colour(color[0], color[1], color[2]));
        else if (color.size() == 4) // RGBA
            vis->SetColour(G4Colour(color[0], color[1], color[2], color[3]));
    }
    catch (DBNotFoundError &e) {
        Log::Die("GeoSc48SourceFactory: DBNotFoundError. Table " + e.table + ", index " + e.index + ", field " + e.field + ".");
    };

    logicalVolume->SetVisAttributes(vis);
}

std::vector<double> GeoSc48SourceFactory::MultiplyVectorByUnit(std::vector<double> v, const double unit) {
    transform(v.begin(), v.end(), v.begin(), bind2nd(std::multiplies<double>(), unit));
    return v;
}


G4PVPlacement* GeoSc48SourceFactory::G4PVPlacementWithCheck(G4Transform3D &Transform3D,
                                          G4LogicalVolume *pCurrentLogical,
                                          const G4String& pName,
                                          G4LogicalVolume *pMotherLogical,
                                          G4bool pMany,
                                          G4int  pCopyNo,
                                          G4bool pSurfChk)
{
    G4PVPlacement* placement = new G4PVPlacement(Transform3D, pCurrentLogical, pName, pMotherLogical, pMany, pCopyNo, false);
    Log::Assert(!pSurfChk || !placement->CheckOverlaps(),
                "GeoSc48SourceFactory: Overlap detected when placing volume " + pName + ". See log for details.");
    return placement;
}

void GeoSc48SourceFactory::Construct(DBLinkPtr table,
                                     const bool checkOverlaps)
{

    // some useful values

    G4RotationMatrix *noRotation = new G4RotationMatrix();
    const bool pMany = false;
    const int pCopyNo = 0;

    try { // to catch DBNotFoundError

        // Preparations
        // ============

        // table index, also serves as prefix for volume names
        const std::string index = table->GetIndex();
        const std::string prefix = index + "_";
        std::pair<G4ThreeVector, G4RotationMatrix*> translation = LoadTranslation( table );

        G4RotationMatrix* Rotation = new G4RotationMatrix();
        if (translation.second != NULL) {Rotation = translation.second;}

        // Find mother volume

        const std::string motherName = table->GetS("mother");

        G4LogicalVolume * const motherLog = Detector::FindLogicalVolume(motherName);
        Log::Assert(motherLog != NULL,
                    "GeoSc48SourceFactory: Unable to find mother volume '" + motherName + "' for '" + index + "'.");

        Log::Assert(Detector::FindPhysicalVolume(motherName) != NULL,
                    "GeoSc48SourceFactory: Mother physical volume '" + motherName + "' not found.");


        // Read parameters from database
        // =============================

        // check for overlap when placing volumes?
        const bool pSurfChk = checkOverlaps;

        // Titanium cylinder (tube)
        const double sampleHeight = table->GetD("sample_height") * mm;
        const double sampleRadius = table->GetD("sample_radius") * mm;
        G4Material * const sampleMaterial = G4Material::GetMaterial(table->GetS("sample_material"));

        // Inner container (tube)
        const double innerContainerHeight = table->GetD("inner_container_height") * mm;
        const double innerContainerRadius = table->GetD("inner_container_radius") * mm;
        G4Material * const innerContainerMaterial = G4Material::GetMaterial(table->GetS("inner_container_material"));

        // Outer container (polycone)
        std::vector<double> outerContainerZ = MultiplyVectorByUnit(table->GetDArray("outer_container_z"), mm);
        std::vector<double> outerContainerRin = MultiplyVectorByUnit(table->GetDArray("outer_container_rin"), mm);
        std::vector<double> outerContainerRout = MultiplyVectorByUnit(table->GetDArray("outer_container_rout"), mm);
        const double outerContainerBottomThickness = table->GetD("outer_container_bottom_thickness") * mm;
        G4Material * const outerContainerMaterial = G4Material::GetMaterial(table->GetS("outer_container_material"));

        // Holder
        const bool holderEnable = table->GetI("holder_enable");
        std::vector<double> holderZ = MultiplyVectorByUnit(table->GetDArray("holder_z"), mm);
        std::vector<double> holderRin = MultiplyVectorByUnit(table->GetDArray("holder_rin"), mm);
        std::vector<double> holderRout = MultiplyVectorByUnit(table->GetDArray("holder_rout"), mm);
        G4Material * const holderMaterial = G4Material::GetMaterial(table->GetS("holder_material"));

        // Weight cylinder
        const bool weightCylinderEnable = table->GetI("weight_cylinder_enable");
        const double weightCylinderHeight = table->GetD("weight_cylinder_height") * mm;
        const double weightCylinderRadius = table->GetD("weight_cylinder_radius") * mm;
        G4Material * const weightCylinderMaterial = G4Material::GetMaterial(table->GetS("weight_cylinder_material"));

        // Screws and cross dowels (connecting outer container and holder)
        const bool screwHolesEnable = table->GetI("screw_holes_enable");
        const bool screwsEnable = table->GetI("screws_enable");
        const int numberOfScrews = table->GetI("number_of_screws");

        const double screwHeadHeight = table->GetD("screw_head_height") * mm;
        const double screwHeadDiameter = table->GetD("screw_head_diameter") * mm;
        const double screwBoltDiameter = table->GetD("screw_bolt_diameter") * mm;
        const double screwDistanceToCenter = table->GetD("screw_distance_to_center") * mm;
        G4Material * const screwMaterial = G4Material::GetMaterial(table->GetS("screw_material"));

        const double crossDowelLength = table->GetD("cross_dowel_length") * mm;
        const double crossDowelDiameter = table->GetD("cross_dowel_diameter") * mm;
        const double crossDowelBoreholePosition = table->GetD("cross_dowel_borehole_position") * mm;
        const double crossDowelDistanceToTop = table->GetD("cross_dowel_distance_to_top") * mm;
        G4Material * const crossDowelMaterial = G4Material::GetMaterial(table->GetS("cross_dowel_material"));

        // sanity checks

        Log::Assert(!(!screwHolesEnable && screwsEnable),
                    "GeoSc48SourceFactory: screw_holes_enable required for screws_enable.");
        Log::Assert(!(!holderEnable && screwsEnable),
                    "GeoSc48SourceFactory: screws_enable is not sensible without holder_enable.");
        Log::Assert(!(!holderEnable && weightCylinderEnable),
                    "GeoSc48SourceFactory: weight_cylinder_enable requires holder_enable.");

        Log::Assert(screwHeadDiameter > screwBoltDiameter,
                    "GeoSc48SourceFactory: screw_head_diameter <= screw_bolt_diameter.");


        // some useful dimensions
        //
        // for the placement of the screws and dowels it is assumed that the
        // "connectors" on both parts correspond to the last/first segment of the
        // polycone

        const double holderConnectorHeight = *(holderZ.begin()+1) - *(holderZ.begin());

        const double crossDowelZ = *(outerContainerZ.end()-1) - crossDowelDistanceToTop;
        const double screwBoltLength = crossDowelDistanceToTop + holderConnectorHeight;

        const double outerContainerZmin = *min_element(outerContainerZ.begin(), outerContainerZ.end());
        const double outerContainerZmax = *max_element(outerContainerZ.begin(), outerContainerZ.end());

        const double holderZmin = *min_element(holderZ.begin(), holderZ.end());
        const double holderZmax = *max_element(holderZ.begin(), holderZ.end());


        // Build the geometry
        // ==================

        // Outer Container

        G4VSolid *outerContainerSolid =
            new G4Polycone(prefix + "outer_container_solid",
                           0.0, twopi,
                           outerContainerZ.size(),
                           &outerContainerZ[0],
                           &outerContainerRin[0],
                           &outerContainerRout[0]);

        if (screwHolesEnable) {

            // prepare hole for the screw
            G4VSolid *holeForScrewInOuterContainerSolid =
                new G4Tubs(prefix + "hole_for_screw_in_outer_container_solid",
                           0.0, screwBoltDiameter / 2.0,
                           (screwBoltLength + screwHeadHeight) / 2.0,
                           0.0, twopi);

            G4ThreeVector holeForScrewInOuterContainerTranslation(screwDistanceToCenter, 0, crossDowelZ + screwBoltLength / 2.0 + screwHeadHeight / 2.0);

            // prepare hole for cross dowel
            // the hole's end is determined by the end of the cross dowel, and it has to be drilled out to the radius at the dowel position
            const double holeForCrossDowelLength = *(outerContainerRout.end()-1) - screwDistanceToCenter + crossDowelBoreholePosition;
            G4ThreeVector holeForCrossDowelTranslation(*(outerContainerRout.end()-1) - holeForCrossDowelLength / 2.0, 0, crossDowelZ);

            G4RotationMatrix *holeForCrossDowelRotation = new G4RotationMatrix();
            holeForCrossDowelRotation->rotateY(twopi/4.0);

            G4VSolid *holeForCrossDowelSolid =
                new G4Tubs(prefix + "hole_for_cross_dowel_solid",
                           0.0, crossDowelDiameter / 2.0,
                           holeForCrossDowelLength / 2.0,
                           0.0, twopi);

            // rotate and drill the holes
            for (int i = 0; i < numberOfScrews; i++) {

                holeForScrewInOuterContainerTranslation = holeForScrewInOuterContainerTranslation.rotateZ(twopi/numberOfScrews);

                outerContainerSolid =
                    new G4SubtractionSolid(prefix + "outer_container_solid",
                                           outerContainerSolid,
                                           holeForScrewInOuterContainerSolid,
                                           noRotation,
                                           holeForScrewInOuterContainerTranslation);

                holeForCrossDowelTranslation = holeForCrossDowelTranslation.rotateZ(twopi/numberOfScrews);
                holeForCrossDowelRotation->rotateZ(twopi/numberOfScrews);
                G4Transform3D holeForCrossDowelTransformation(*holeForCrossDowelRotation, holeForCrossDowelTranslation);

                outerContainerSolid =
                    new G4SubtractionSolid(prefix + "OuterContainer",
                                           outerContainerSolid,
                                           holeForCrossDowelSolid,
                                           holeForCrossDowelTransformation);
            }
        }

        G4LogicalVolume *outerContainerLog = new G4LogicalVolume(outerContainerSolid, outerContainerMaterial, prefix + "outer_container_logical");
        SetColor(table, "outer_container_vis_color", outerContainerLog);

        G4ThreeVector outerContainerPlacementPosition(translation.first.x(), translation.first.y(),
            translation.first.z() - outerContainerZmin - outerContainerBottomThickness - innerContainerHeight / 2.0);

        G4Transform3D outerContainerPlacementTransformation(*Rotation, outerContainerPlacementPosition);

        G4PVPlacementWithCheck(
            outerContainerPlacementTransformation,
            outerContainerLog,
            outerContainerLog->GetName(),
            motherLog,
            pMany, pCopyNo, pSurfChk);

        // Inner Container

        G4VSolid *innerContainerSolid =
            new G4Tubs(prefix + "inner_container_solid",
                       0.0, innerContainerRadius,
                       innerContainerHeight / 2.0,
                       0.0, twopi);

        G4LogicalVolume *innerContainerLog = new G4LogicalVolume(innerContainerSolid, innerContainerMaterial, prefix + "inner_container_logical");
        SetColor(table, "inner_container_vis_color", innerContainerLog);

        G4ThreeVector innerContainerPlacementPosition(0.0, 0.0,
            outerContainerZmin + innerContainerHeight / 2.0 + outerContainerBottomThickness);

        G4Transform3D innerContainerPlacementTransformation(*noRotation, innerContainerPlacementPosition);

        G4PVPlacementWithCheck(
            innerContainerPlacementTransformation,
            innerContainerLog,
            innerContainerLog->GetName(),
            outerContainerLog,
            pMany, pCopyNo, pSurfChk);


        // Sample

        G4VSolid *sampleSolid =
            new G4Tubs(prefix + "sample_solid",
                       0.0, sampleRadius,
                       sampleHeight / 2.0,
                       0.0, twopi);

        G4LogicalVolume *sampleLog = new G4LogicalVolume(sampleSolid, sampleMaterial, prefix + "sample_logical");
        SetColor(table, "sample_vis_color", sampleLog);

        G4ThreeVector samplePlacementPosition(0.0, 0.0, 0.0);
        G4Transform3D samplePlacementTransformation(*noRotation, samplePlacementPosition);

        G4PVPlacementWithCheck(
            samplePlacementTransformation,
            sampleLog,
            sampleLog->GetName(),
            innerContainerLog,
            pMany, pCopyNo, pSurfChk);

        // Holder (if enabled)

        G4ThreeVector holderPlacementPosition(0,0,0);

	G4ThreeVector holderPosition = G4ThreeVector(0.0, 0.0, outerContainerZmax - holderZmin);

	holderPlacementPosition = outerContainerPlacementPosition + *Rotation*holderPosition;

        if (holderEnable) {
            G4VSolid *holderSolid =
                new G4Polycone(prefix + "holder_solid",
                               0.0, twopi,
                               holderZ.size(),
                               &holderZ[0],
                               &holderRin[0],
                               &holderRout[0]);

            // drill holes, if enabled
            if (screwHolesEnable) {

                G4VSolid *holeForScrewInHolderSolid =
                    new G4Tubs(prefix + "hole_for_screw_in_holder_solid",
                               0.0, screwBoltDiameter / 2.0,
                               screwBoltLength / 2.0,
                               0.0, twopi);

                G4ThreeVector holeForScrewInHolderTranslation(screwDistanceToCenter, 0,
                    holderZmin - crossDowelDistanceToTop + screwBoltLength/2.0);

                // rotate and drill
                for (int i = 0; i < numberOfScrews; i++) {
                    holeForScrewInHolderTranslation = holeForScrewInHolderTranslation.rotateZ(twopi/numberOfScrews);
                    G4Transform3D holeForScrewInHolderTransformation(*noRotation, holeForScrewInHolderTranslation);

                    holderSolid =
                        new G4SubtractionSolid(prefix + "holder_solid",
                                    holderSolid,
                                    holeForScrewInHolderSolid,
                                    holeForScrewInHolderTransformation);
                }
            }

            G4LogicalVolume *holderLog = new G4LogicalVolume(holderSolid, holderMaterial, prefix + "holder_logical");
            SetColor(table, "holder_vis_color", holderLog);


            G4Transform3D holderPlacementTransformation(*Rotation, holderPlacementPosition);

            G4PVPlacementWithCheck(
                holderPlacementTransformation,
                holderLog,
                holderLog->GetName(),
                motherLog,
                pMany, pCopyNo, pSurfChk);
        }

        // place screws and dowels, if requested
        if (screwsEnable) {

            // assemble and place screws

            G4VSolid *screwBoltSolid=
                    new G4Tubs(prefix + "screw_bolt_solid",
                               0.0,
                               screwBoltDiameter / 2.0,
                               (screwBoltLength + screwHeadHeight) / 2.0,
                               0.0, twopi);

            G4VSolid *screwHeadSolid =
                    new G4Tubs(prefix + "screw_head_solid",
                               0.0,
                               screwHeadDiameter / 2.0,
                               screwHeadHeight / 2.0,
                               0.0, twopi);

            G4UnionSolid *screwSolid =
                new G4UnionSolid(prefix + "screw_solid",
                                 screwHeadSolid,
                                 screwBoltSolid,
                                 noRotation,
                                 G4ThreeVector(0, 0, - screwBoltLength / 2.0));

            G4LogicalVolume *screwLog = new G4LogicalVolume(screwSolid, screwMaterial, prefix + "screw_logical");
            SetColor(table, "screw_vis_color", screwLog);

            G4ThreeVector screwPlacementPosition(screwDistanceToCenter, 0,
                holderPosition.z() + holderZmin + screwBoltLength - crossDowelDistanceToTop + screwHeadHeight / 2.0);

            // rotate and place screws
            for (int i = 0; i < numberOfScrews; i++) {
                screwPlacementPosition = screwPlacementPosition.rotateZ(twopi/numberOfScrews);
                G4Transform3D screwPlacementTransformation(*Rotation, *Rotation*screwPlacementPosition + outerContainerPlacementPosition);

                G4PVPlacementWithCheck(
                    screwPlacementTransformation,
                    screwLog,
                    screwLog->GetName() + "_" + to_string(i),
                    motherLog,
                    pMany, pCopyNo, pSurfChk);
            }

            // assemble and place dowels

            G4VSolid *crossDowelSolid =
                    new G4Tubs(prefix + "cross_dowel_solid",
                               0.0, crossDowelDiameter / 2.0,
                               crossDowelLength / 2.0,
                               0.0, twopi);

            G4VSolid *crossDowelBoreholeSolid =
                    new G4Tubs(prefix + "cross_dowel_borehole",
                               0.0,
                               screwBoltDiameter / 2.0,
                               crossDowelDiameter / 2.0,
                               0.0, twopi);

            G4RotationMatrix *crossDowelBoreholeRotation = new G4RotationMatrix();
            crossDowelBoreholeRotation->rotateY(twopi/4.0); // place hole horizontally
            G4ThreeVector crossDowelBoreholeTranslation(0, 0, -crossDowelLength/2.0 + crossDowelBoreholePosition);

            crossDowelSolid = new G4SubtractionSolid(prefix + "cross_dowel_solid",
                                                     crossDowelSolid,
                                                     crossDowelBoreholeSolid,
                                                     crossDowelBoreholeRotation,
                                                     crossDowelBoreholeTranslation);

            G4LogicalVolume *crossDowelLog = new G4LogicalVolume(crossDowelSolid, crossDowelMaterial, prefix + "cross_dowel_logical");
            SetColor(table, "cross_dowel_vis_color", crossDowelLog);

            G4RotationMatrix *crossDowelRotation = new G4RotationMatrix();
            crossDowelRotation->rotateY(twopi/4.0); // dowels are horizontal

            G4ThreeVector crossDowelPlacementPosition(
                screwDistanceToCenter + crossDowelLength / 2.0 - crossDowelBoreholePosition,
                0, holderPosition.z() + holderZmin - crossDowelDistanceToTop);

            // rotate and place dowels
            for (int i = 0; i < numberOfScrews; i++) {
                crossDowelPlacementPosition = crossDowelPlacementPosition.rotateZ(twopi/numberOfScrews);
                crossDowelRotation->rotateZ(twopi/numberOfScrews);
                G4Transform3D crossDowelPlacementTransformation(*crossDowelRotation, crossDowelPlacementPosition);

                G4PVPlacementWithCheck(
                    crossDowelPlacementTransformation,
                    crossDowelLog,
                    crossDowelLog->GetName() + "_" + to_string(i),
                    outerContainerLog,
                    pMany, pCopyNo, pSurfChk);
            }
        }


        // Weight Cylinder

        if (weightCylinderEnable) {

            G4VSolid *weightCylinderSolid =
                new G4Tubs(prefix + "weight_cylinder_solid",
                           0.0, weightCylinderRadius,
                           weightCylinderHeight / 2.0,
                           0.0, twopi);

            G4LogicalVolume *weightCylinderLog = new G4LogicalVolume(weightCylinderSolid, weightCylinderMaterial, prefix + "weight_cylinder_logical");
            SetColor(table, "weight_cylinder_vis_color", weightCylinderLog);

            G4ThreeVector weightCylinderPlacementPosition(holderPosition.x(),
                holderPosition.y(),
                holderPosition.z() + holderZmax + weightCylinderHeight/2.0);
            G4Transform3D weightCylinderTransformation(*Rotation, *Rotation*weightCylinderPlacementPosition + outerContainerPlacementPosition);

            G4PVPlacementWithCheck(
                weightCylinderTransformation,
                weightCylinderLog,
                weightCylinderLog->GetName(),
                motherLog,
                pMany, pCopyNo, pSurfChk);
        }
    }
    catch(DBNotFoundError &e) {
        Log::Die("GeoSc48SourceFactory: DBNotFoundError. Table " + e.table + ", index " + e.index + ", field " + e.field + ".");
    };
    return;
}

} // namespace RAT