#===----------------------------------------------------------------------===#
#
#         STAIRLab -- STructural Artificial Intelligence Laboratory
#
#===----------------------------------------------------------------------===#
#
import os.path
import shutil
from typing import NewType
import sys, json
import zipfile
from pathlib import Path
from abc import abstractmethod
import contextlib

class Event: pass

RunID = NewType("RunID", int)

MetricType = NewType("MetricType", str)

from .utilities import read_model
from .metrics import (
     accel_response_history_plot,
     column_strain_state_metric,
     peak_acceleration_metric,
     peak_drift_metric
)

OPENSEES = [
    sys.executable, "-m", "opensees",
]


@contextlib.contextmanager
def new_cd(x):
    d = os.getcwd()

    # This could raise an exception, but it's probably
    # best to let it propagate and let the caller
    # deal with it, since they requested x
    os.chdir(x)

    try:
        yield

    finally:
        # This could also raise an exception, but you *really*
        # aren't equipped to figure out what went wrong if the
        # old working directory can't be restored.
        os.chdir(d)

class classproperty(property):
    def __get__(self, cls, owner):
        return classmethod(self.fget).__get__(None, owner)()


class Runner:
    def __init__(self, conf: dict):

        if isinstance(conf, dict):
            # Create from dict when posted from API; this
            # is used to create a new PredictorModel
            self.name: str   = conf["name"]
            self.description = conf.get("description", "")
            self.conf        = conf["config"]
            self.metrics     = conf["metrics"]
            self.entry_point = conf["entry_point"]
            self.active = conf.get("active", True)
        else:
            # Create from PredictorModel when loaded from database.
            # This is done when running analysis
            self.name: str = conf.name
            self.description = "" # conf.description
            self.conf = conf.config
            self.entry_point = conf.entry_point
            self.metrics = conf.metrics
            self.active = conf.active
            try:
                self.model_file = Path(conf.config_file.path).resolve()
            except ValueError:
                pass

            self.out_dir = Path(__file__).parents[0]/"Predictions"
            self.runs = {}

    @abstractmethod
    def newPrediction(self, event)->RunID: ...

    @abstractmethod
    def runPrediction(self, run_id)->bool: ...

    def getMetricList(self)->list:
        return self.metrics

    def activateMetric(self, type, rid=None)->bool:
        return False

    @abstractmethod
    def getMetricData(self, run: RunID, metric: MetricType)->dict: ...


class OpenSeesRunner(Runner):
    @property
    def platform(self):
        return self.conf.get("platform", "")


    @classmethod
    def create(cls, asset, request, form):
        predictor = form.save(commit=False)

        predictor.entry_point = [
            sys.executable, "-m", "opensees"
        ]
        predictor.config = {}
        predictor.protocol = "BRACE2_CLI_PREDICTOR_T4"
        predictor.active = True
        # predictor.metrics = self.getMetricList()
        return predictor


    @classproperty
    def schema(cls):
        from .opensees import schemas
        return {
            "title": "Structural Model",
            "options": {"disable_collaps": True},
            "schema": "http://json-schema.org/draft-04/schema#",
            "type": "object",
            "properties": {
                "platform": {
                  "type": "string",
                  "title": "Platform",
                  "enum": ["OpenSees","CSiBridge"]
                },
                "model":    schemas.load("hwd_conf.schema.json"),
                "analysis": schemas.load("hwd_analysis.schema.json"),
            }
        }

    def getMetricList(self):
        return [
            "COLUMN_STRAIN_STATES",
            "PEAK_ACCEL",
            "PEAK_DRIFT",
            # "ACC_RESPONSE_HISTORY",
        ]

    def newPrediction(self, event, output_directory = None):
        """
        Create a new prediction run and return the run_id. If output_directory is None,
        the output directory will be created automatically. Otherwise, the output directory
        will be copied to the new output directory.
        """
        event = event.event_file.path
        if output_directory is not None:
            # this case will eventually be deleted, its just for
            # debugging metric renderers.
            run_id = "0"
            self.runs[run_id] = {
                "run_output_directory": Path(output_directory)
            }

        else:
            # Calculate next output directory and
            # create directory if it doesn't exist
            out_dir = self.out_dir
            if not out_dir.is_dir():
                (out_dir/"0").mkdir(parents=True)

            latestDir = list(sorted((f for f in out_dir.iterdir() if f.is_dir()), key=lambda m: int(m.name)))[-1]
            run_id = int(latestDir.name)+1
            run_dir = out_dir/str(run_id)
            run_dir = run_dir.resolve()
            run_dir.mkdir(parents=True, exist_ok=False)

            # Copy files to run directory
            shutil.copyfile(event, run_dir/"event.zip")
            shutil.copyfile(self.model_file.resolve(), run_dir/self.model_file.name)

            if self.model_file.suffix == ".zip":
                with zipfile.ZipFile(self.model_file, 'r') as zip_ref:
                    zip_ref.extractall(run_dir)
                model_file = (run_dir/"nonlinear.tcl").resolve()

            elif self.model_file.suffix == ".b2k":
                pass

            elif self.model_file.suffix == ".tcl":
                model_file = (run_dir/self.model_file.name).resolve()

            self.runs[run_id] = {
                "run_output_directory": run_dir,
                "event_file_name": Path(event),
                "model_file": model_file,
                **self.conf
            }

            with open(out_dir/str(run_id)/"conf.json", "w") as f:
                json.dump({k: str(v) for k,v in self.runs[run_id].items()}, f)

        return run_id


    def _load_config(self, run_id):
        run_dir =  self.out_dir/str(run_id)
        with open(run_dir/"conf.json","r") as f:
            self.runs[run_id] = json.load(f)

        self.model_file = Path(self.runs[run_id]["model_file"])


    def runPrediction(self, run_id, scale: float = None):
        if run_id not in self.runs:
            self._load_config(run_id)

        event_file_path = os.path.relpath(self.runs[run_id]["event_file_name"],
                                          self.model_file.parents[0])
        output_directory = os.path.relpath(self.runs[run_id]["run_output_directory"],
                                           self.model_file.parents[0])

        event_file_path = self.runs[run_id]["event_file_name"]

        # Create model
        import opensees.openseespy as ops

        import sys
        model = ops.Model(echo_file=sys.stdout)
        model.eval("set argv {}")
        with new_cd(self.runs[run_id]["run_output_directory"]):
            model.eval(f"source {self.runs[run_id]['model_file']}")

            model.eval(f"print -json -file modelDetails.json")

            model.eval(f"set python {sys.executable}")

            model.eval(r"""
                proc py {args} {
                    global python
                    eval "[exec {*}$python {*}$args]"
                }
                
                proc pt {args} {
                    global python
                    puts "[exec {*}$python {*}$args]"
                }

                proc write_modes {mode_file nmodes} {
                    set fid_modes [open $mode_file w+]
                    for {set m 1} {$m <= $nmodes} {incr m} {
                        puts $fid_modes "$m:"
                        foreach n [getNodeTags] {
                        puts $fid_modes "  $n: \[[join [nodeEigenvector $n $m] {, }]\]";
                        }
                    }
                    close $fid_modes
                }
                proc write_displacements {file_name {resp Disp}} {
                    set fid [open "$file_name" "w+"]
                    puts $fid "[getTime]:"
                    foreach n [getNodeTags] {
                        puts $fid "    $n: \[[join [node${resp} $n] {, }]\]";
                    }
                    close $fid;
                }
            """)

            #
            # Run gravity analysis
            #
            model.eval("""
            wipeAnalysis
            test NormDispIncr 1.0e-8 10 0;
            algorithm Newton;
            integrator LoadControl 0.1;
            numberer Plain;
            constraints Transformation;
            system SparseGeneral;
            analysis Static;
            analyze 10;
            #  write_displacements "dispsGrav.yaml"
            """)

            #
            # DAMPING
            #
            model.eval(r"""
            set nmodes 8; # Number of modes to analyze for modal analysis

            # set wb [eigen -fullGenLapack $nmodes];
            # puts "\tFundamental-Period After Gravity Analysis:"
            # for {set iPd 1} {$iPd <= $nmodes} {incr iPd 1} {
            #     set wwb [lindex $wb $iPd-1];
            #     set Tb [expr 2*$pi/sqrt($wwb)];
            #     puts "\tPeriod$iPd= $Tb"
            # }
            # write_modes $output_directory/modesPostG.yaml $nmodes
            # remove recorders

            set nmodes [tcl::mathfunc::max {*}$damping_modes $nmodes]
            set lambdaN [eigen  -fullGenLapack $nmodes];

            # set lambdaN [eigen $nmodes];
            if {$damping_type == "rayleigh"} {
                set nEigenI [lindex $damping_modes 0];                  # first rayleigh damping mode
                set nEigenJ [lindex $damping_modes 1];                  # second rayleigh damping mode
                set iDamp   [lindex $damping_ratios 0];                 # first rayleigh damping ratio
                set jDamp   [lindex $damping_ratios 1];                 # second rayleigh damping ratio
                set lambdaI [lindex $lambdaN [expr $nEigenI-1]];
                set lambdaJ [lindex $lambdaN [expr $nEigenJ-1]];
                set omegaI [expr $lambdaI**0.5];
                set omegaJ [expr $lambdaJ**0.5];
                set TI [expr 2.0*$pi/$omegaI];
                set TJ [expr 2.0*$pi/$omegaJ];
                set alpha0 [expr 2.0*($iDamp/$omegaI-$jDamp/$omegaJ)/(1/$omegaI**2-1/$omegaJ**2)];
                set alpha1 [expr 2.0*$iDamp/$omegaI-$alpha0/$omegaI**2];
                puts "\tRayleigh damping parameters:"
                puts "\tmodes: $nEigenI, $nEigenJ ; ratios: $iDamp, $jDamp"
                puts "\tTI = $TI; TJ = $TJ"
                puts "\tlambdaI = $lambdaI; lambdaJ = $lambdaJ"
                puts "\tomegaI = $omegaI; omegaJ = $omegaJ"
                puts "\talpha0 = $alpha0; alpha1 = $alpha1"
                rayleigh $alpha0 0.0 0.0 $alpha1;

            } elseif {$damping_type == "modal"} {
                # needs a bit of edit. currently assuming that the ratios are applied in order at the first modes. but should be applied at the specified damping_modes modes.
                set nratios [llength $damping_ratios]
                puts "\tModal damping parameters:"
                puts "\tratios of $damping_ratios at the first $nratios modes"
                for {set i 1} {$i <= [expr $nmodes - $nratios]} {incr i} {
                    lappend damping_ratios 0
                }
                modalDamping {*}$damping_ratios
            }
            """)


            #
            # DYNAMIC RECORDERS
            #

            ## COLUMN SECTION DEFORMATIONS AT TOP AND BOTTOM FOR STRAIN-BASED DAMAGE STATES
            column_strains = tuple(k["key"] for k in CONFIG["columns"] if k["strain"])
            if len(column_strains) > 0:
                model.recorder("Element",  "section", 1, "deformation", xml="eleDef1.txt", ele=column_strains) # section 1 deformation]
                model.recorder("Element",  "section", 4, "deformation", xml="eleDef4.txt", ele=column_strains) # section 4 deformation]



            #
            # Run dynamic analysis
            #
            model.eval(f"""
            wipeAnalysis
            # Uniform Support Excitation
#           lassign [pt -m  CE58658.makePattern {event_file_path} --scale $dynamic_scale_factor --node $input_location] dt steps
#           lassign [py -m  CE58658.makePattern {event_file_path} --scale $dynamic_scale_factor --node $input_location] dt steps
            set dt 0.1
            set steps 3
            """)

            # RESPONSE HISTORY RECORDERS

            model.recorder("Node", "accel", xml="model/AA_all.txt", timeSeries=(1, 2), dof=(1, 2))
            model.recorder("Node", "accel", xml="model/RD_all.txt", dof=(1, 2))

            column_nodes = tuple(k["node"] for k in CONFIG["bents"] if k["record"])
            model.recorder("Node", "accel", file="TopColAccel_X_txt.txt", timeSeries=1 , node=column_nodes, dof=1)
            model.recorder("Node", "accel", file="TopColAccel_Y_txt.txt", timeSeries=2 , node=column_nodes, dof=2)
            model.recorder("Node", "disp",  file="TopColDrift_X_txt.txt", node=column_nodes, dof=1)
            model.recorder("Node", "disp",  file="TopColDrift_Y_txt.txt", node=column_nodes, dof=2)

            model.eval("""
            set dtfact 1;
            set Tol                1.0e-8;
            set maxNumIter        100;
            set printFlag        0;
            set TestType        EnergyIncr;
            set NewmarkGamma    0.50;
            set NewmarkBeta        0.25;
            constraints Transformation;
            numberer RCM;
            test $TestType $Tol $maxNumIter $printFlag;
            set algorithmType   "Newton";
            system BandSPD;
            integrator Newmark $NewmarkGamma $NewmarkBeta;

            algorithm {*}$algorithmType;
            analysis Transient;

            set DtAnalysis         $dt;
            set TmaxAnalysis     [expr $dt*$steps];
            set Nsteps             $steps;
            if {$dynamic_truncated != 0} {
                set Nsteps             $dynamic_timesteps;
            }
            puts "\tGround Motion: dt= $DtAnalysis, NumPts= $Nsteps, TmaxAnalysis= $TmaxAnalysis";

            puts "\tRunning dynamic ground motion analysis..."
            set t3 [clock clicks -millisec];
            catch {progress create $Nsteps} _

            analyze 2 $DtAnalysis;

#           for {set ik 1} {$ik <= $Nsteps} {incr ik 1} {
#               catch {progress update} _ 
#               set ok      [analyze 1 $DtAnalysis];
#           }
            """)

            model.wipe()


    def getMetricData(self, run_id:int, type:str)->dict:
        import orjson
        def _clean_json(d):
            return orjson.loads(orjson.dumps(d,option=orjson.OPT_SERIALIZE_NUMPY))

        if run_id not in self.runs:
            self._load_config(run_id)

        run_data = self.runs.get(run_id, None)
        config = run_data

        if run_data is not None:
            output_dir = Path(run_data["run_output_directory"])
        else:
            output_dir = self.out_dir/str(run_id)

        # with open(output_dir/"modelDetails.json", "r") as f:
        #     model = json.load(f)

        model = read_model(output_dir/"modelDetails.json")

        if type == "COLUMN_STRAIN_STATES":
            return _clean_json(column_strain_state_metric(model, output_dir, config))

        elif type == "PEAK_ACCEL":
            return _clean_json(peak_acceleration_metric(output_dir, config))

        elif type == "PEAK_DRIFT":
            return _clean_json(peak_drift_metric(model, output_dir, config))

        elif type == "ACC_RESPONSE_HISTORY":
            # config = CONFIG
#           return accel_response_history_plot(output_dir, config)
            return {}
        return {}

