#!/usr/bin/env python3
from pathlib import Path
from enum import Enum
import argparse

import acts
import acts.examples
from acts.examples.odd import getOpenDataDetector

u = acts.UnitConstants


# Graciously taken from https://stackoverflow.com/a/60750535/4280680
class EnumAction(argparse.Action):
    """
    Argparse action for handling Enums
    """

    def __init__(self, **kwargs):
        # Pop off the type value
        enum_type = kwargs.pop("enum", None)

        # Ensure an Enum subclass is provided
        if enum_type is None:
            raise ValueError("type must be assigned an Enum when using EnumAction")
        if not issubclass(enum_type, Enum):
            raise TypeError("type must be an Enum when using EnumAction")

        # Generate choices from the Enum
        kwargs.setdefault("choices", tuple(e.name for e in enum_type))

        super(EnumAction, self).__init__(**kwargs)

        self._enum = enum_type

    def __call__(self, parser, namespace, values, option_string=None):
        for e in self._enum:
            if e.name == values:
                setattr(namespace, self.dest, e)
                break
        else:
            raise ValueError("%s is not a validly enumerated algorithm." % values)


from acts.examples.reconstruction import SeedingAlgorithm


def runSeeding(
    trackingGeometry,
    field,
    outputDir,
    s=None,
    seedingAlgorithm=SeedingAlgorithm.GridTriplet,
):
    from acts.examples.simulation import (
        addParticleGun,
        EtaConfig,
        PhiConfig,
        ParticleConfig,
        addFatras,
        addDigitization,
        ParticleSelectorConfig,
        addDigiParticleSelection,
    )

    s = s or acts.examples.Sequencer(
        events=100, numThreads=-1, logLevel=acts.logging.INFO
    )
    rnd = acts.examples.RandomNumbers(seed=42)
    outputDir = Path(outputDir)

    addParticleGun(
        s,
        EtaConfig(-2.0, 2.0),
        ParticleConfig(4, acts.PdgParticle.eMuon, True),
        PhiConfig(0.0, 360.0 * u.degree),
        multiplicity=2,
        outputDirCsv=outputDir / "csv",
        outputDirRoot=outputDir,
        rnd=rnd,
    )

    addFatras(
        s,
        trackingGeometry,
        field,
        outputDirCsv=outputDir / "csv",
        outputDirRoot=outputDir,
        rnd=rnd,
    )

    srcdir = Path(__file__).resolve().parent.parent.parent.parent
    addDigitization(
        s,
        trackingGeometry,
        field,
        digiConfigFile=srcdir / "Examples/Configs/generic-digi-smearing-config.json",
        rnd=rnd,
    )

    addDigiParticleSelection(
        s,
        ParticleSelectorConfig(
            pt=(1.0 * u.GeV, None),
            eta=(-2.5, 2.5),
            measurements=(9, None),
            removeNeutral=True,
        ),
    )

    from acts.examples.reconstruction import (
        addSeeding,
        SeedFinderConfigArg,
        SeedFinderOptionsArg,
    )

    addSeeding(
        s,
        trackingGeometry,
        field,
        SeedFinderConfigArg(
            r=(None, 200 * u.mm),  # rMin=default, 33mm
            deltaR=(1 * u.mm, 300 * u.mm),
            collisionRegion=(-250 * u.mm, 250 * u.mm),
            z=(-2000 * u.mm, 2000 * u.mm),
            maxSeedsPerSpM=1,
            sigmaScattering=50,
            radLengthPerSeed=0.1,
            minPt=500 * u.MeV,
            impactMax=3 * u.mm,
        ),
        SeedFinderOptionsArg(
            bFieldInZ=2 * u.T,
        ),
        acts.logging.VERBOSE,
        seedingAlgorithm=seedingAlgorithm,
        geoSelectionConfigFile=srcdir / "Examples/Configs/generic-seeding-config.json",
        outputDirRoot=outputDir,
    )
    return s


if "__main__" == __name__:
    p = argparse.ArgumentParser(
        description="Example script to run seed finding",
    )

    p.add_argument(
        "--algorithm",
        action=EnumAction,
        enum=SeedingAlgorithm,
        default=SeedingAlgorithm.GridTriplet,
        help="Select the seeding algorithm to use",
    )

    args = p.parse_args()
    # detector = getOpenDataDetector()
    detector = acts.examples.GenericDetector()
    trackingGeometry = detector.trackingGeometry()

    field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))

    runSeeding(
        trackingGeometry, field, outputDir=Path.cwd(), seedingAlgorithm=args.algorithm
    ).run()