from typing import Dict, Tuple
import logging
import pandas as pd
import numpy as np
import warnings
import attr
from attr.validators import instance_of
from enum import Enum
from pysight.nd_hist_generator.movie import ImagingSoftware
from ..line_signal_validators.scanimage import ScanImageLineValidator
from ..line_signal_validators.mscan import MScanLineValidator

class LinesType(Enum):
    SWEEPS_REBUILD = "sweeps-rebuild"
    SWEEPS_FROM_SCRATCH = "sweeps-from-scratch"
    CORRUPT = "corrupt"
    REBUILD = "rebuild"
    SCANIMAGE = "si_valid"
    MSCAN = "mscan_valid"
    FROM_SCRATCH = "from_scratch"

@attr.s(slots=True)
class SignalValidator:
    """
    Parse a dictionary of all recorded signals and verify the data's integrity

    :param dict dict_of_data: Dictionary with keys corresponding to signal type,
                            and values as DataFrames.
    :param list data_to_grab: Column names from the DF of the data to parse.
    :param int num_of_lines: Number of lines (x pixels) in the frame)
    :param int num_of_frames: Number of frames of the data. If "Frames" is not in the
                            dict_of_data keys then this number is disregarded.
    :param float binwidth: Multiscaler binwidth in seconds.
    :param float line_freq: Line frequency of the fast scanner (x) in Hz.
    :param bool bidir: Whether the scan was bidirectional.
    :param float delay_between_frames: Time between subsequent frames. Depends on imaging
                                        software and hardware.
    :param bool use_sweeps: Use the sweeps of the multiscaler as an indicator of new lines.
    :param float bidir_phase: A number in microseconds typically reported by the scanning software.
                                Helps with alignment of the generated image when using a bidirectional scan.
    :param ImagingSoftware imaging_software: Software used for imaging.
    :param float change_thresh: Percent-like [0, 1] threshold to discard irregular line signals.
                                Above which PySight will terminate with an exception.
    :param dict lst_metadata: Acquisiton parameters as recorded in the multiscaler.
    """

    dict_of_data = attr.ib(validator=instance_of(dict))
    data_to_grab = attr.ib(default=["abs_time", "sweep"], validator=instance_of(list))
    num_of_lines = attr.ib(default=512, validator=instance_of(int))
    num_of_frames = attr.ib(default=1, validator=instance_of(int))
    binwidth = attr.ib(default=800e-12, validator=instance_of(float))
    line_freq = attr.ib(default=7900.0, validator=instance_of(float))
    bidir = attr.ib(default=False, validator=instance_of(bool))
    delay_between_frames = attr.ib(default=0.001_135_5, validator=instance_of(float))
    use_sweeps = attr.ib(default=False, validator=instance_of(bool))
    bidir_phase = attr.ib(default=-2.7, validator=instance_of(float))
    image_soft = attr.ib(
        default=ImagingSoftware.SCANIMAGE.value, validator=instance_of(str)
    )
    change_thresh = attr.ib(default=0.3, validator=instance_of(float))
    lst_metadata = attr.ib(factory=dict, validator=instance_of(dict))
    handle_line_cases = attr.ib(init=False)
    last_event_time = attr.ib(init=False)
    line_type_data = attr.ib(init=False)
    line_delta = attr.ib(init=False)

    def __attrs_post_init__(self):
        self.handle_line_cases = {
            LinesType.CORRUPT: self.__corrupt,
            LinesType.FROM_SCRATCH: self.__from_scratch,
            LinesType.SWEEPS_FROM_SCRATCH: self.__sweeps_from_scratch,
            LinesType.REBUILD: self.__rebuild,
            LinesType.SWEEPS_REBUILD: self.__sweeps_rebuild,
            LinesType.SCANIMAGE: self.__si,
            LinesType.MSCAN: self.__mscan,
        }

    @property
    def frame_delay(self):
        return np.uint64(self.delay_between_frames / self.binwidth)

    @property
    def total_sweep_time(self):
        return (
            int(self.lst_metadata["fstchan"])
            + int(self.lst_metadata["data_range"])
            + int(self.lst_metadata["holdafter"])
        )

    @property
    def max_sweep(self):
        return int(self.lst_metadata["swpreset"])

    def run(self):
        """ Main pipeline """
        if "Frames" in self.dict_of_data:
            self.num_of_frames = (
                self.dict_of_data["Frames"].shape[0] + 1
            )  # account for first frame
        self.last_event_time = self.__calc_last_event_time()
        self.line_type_data: LinesType = self.__match_line_data_to_case(
            lines=self.dict_of_data.get("Lines")
        )
        self.dict_of_data, self.line_delta = self.handle_line_cases[
            self.line_type_data
        ]()
        self.dict_of_data = self.__validate_frame_input()
        try:
            self.dict_of_data["Laser"] = self.__validate_laser_input(
                self.dict_of_data["Laser"]
            )
        except KeyError:
            pass

        self.__validate_created_data_channels()

    def __bins_bet_lines(self) -> int:
        """
        Calculate the line difference in bins, according to the frequency
        of the line signal.
        :param line_freq: Frequency of lines in Hz. Regardless of bidirectional scanning.
        :return: int - line delta
        """
        freq_in_bins = 1 / (self.line_freq * self.binwidth)
        return int(freq_in_bins / 2) if self.bidir else freq_in_bins

    def __calc_last_event_time(self) -> int:
        """
        Find the last event time for the experiment. Logic as follows:
        No lines / frames data given: Last event time is the last photon time.
        Only lines data given: The last start-of-frame time is created, and the difference between subsequent frames
        in the data is added.
        Frames data exists: The last frame time plus the difference between subsequent frames is the last event time.
        :return int: Last event time
        """

        if "Frames" in self.dict_of_data:
            last_frame_time = self.dict_of_data["Frames"].loc[:, "abs_time"].iloc[-1]

            if self.dict_of_data["Frames"].shape[0] == 1:
                return int(2 * last_frame_time)
            else:
                frame_diff = int(
                    self.dict_of_data["Frames"].loc[:, "abs_time"].diff().mean()
                )
                return int(last_frame_time + frame_diff)

        if "Lines" in self.dict_of_data:
            num_of_lines_recorded = self.dict_of_data["Lines"].shape[0]
            div, mod = divmod(num_of_lines_recorded, self.num_of_lines)

            # if num_of_lines_recorded > (self.num_of_lines * div):  # excessive number of lines
            #     last_line_of_last_frame = self.dict_of_data['Lines'].loc[:, 'abs_time']\
            #         .iloc[div * self.num_of_lines - 1]
            #     line_diff = int(self.dict_of_data['Lines'].iloc[:, 0].diff().mean())
            #     next_frame_event = self.dict_of_data['Lines'].iloc[div * self.num_of_lines, :]
            #     next_frame_event['abs_time'] += np.uint64(10 * line_diff)
            #     self.dict_of_data['Lines'] = self.dict_of_data['Lines'].loc[
            #                                  self.dict_of_data['Lines'].loc[:, 'abs_time'] <= last_line_of_last_frame,:]
            #     if self.image_soft == ImagingSoftware.SCANIMAGE.value:
            #         self.dict_of_data['Lines'] = self.dict_of_data['Lines'].append(next_frame_event)
            #     return int(last_line_of_last_frame + line_diff)

            if (
                mod == 0
            ):  # number of lines contained exactly in number of lines per frame
                return int(
                    self.dict_of_data["Lines"].loc[:, "abs_time"].iloc[-1]
                    + self.dict_of_data["Lines"].loc[:, "abs_time"].diff().mean()
                )

            elif (num_of_lines_recorded < self.num_of_lines * (div + 1)) and (div > 0):
                last_line_of_last_frame = (
                    self.dict_of_data["Lines"]
                    .loc[:, "abs_time"]
                    .iloc[div * self.num_of_lines - 1]
                )
                line_diff = int(self.dict_of_data["Lines"].iloc[:, 0].diff().mean())
                # missing_lines = self.num_of_lines - mod
                # next_frame_event = self.dict_of_data['Lines'].iloc[div * self.num_of_lines, :]
                # next_frame_event['abs_time'] += np.uint64(10 * line_diff)
                # if self.image_soft == ImagingSoftware.SCANIMAGE.value:
                #     self.dict_of_data['Lines'] = self.dict_of_data['Lines'].append(next_frame_event)
                #
                return int(last_line_of_last_frame + line_diff)

            elif div == 0:
                line_diff = int(self.dict_of_data["Lines"].iloc[:, 0].diff().mean())
                missing_lines = self.num_of_lines - num_of_lines_recorded + 1
                return self.dict_of_data["Lines"].iloc[-1, 0] + (
                    line_diff * missing_lines
                )

        # Just PMT data
        max_pmt1 = self.dict_of_data["PMT1"].loc[:, "abs_time"].max()
        try:
            max_pmt2 = self.dict_of_data["PMT2"].loc[:, "abs_time"].max()
        except KeyError:
            return max_pmt1
        else:
            return max(max_pmt1, max_pmt2)

    def __match_line_data_to_case(self, lines: pd.Series) -> LinesType:
        """
        Find the suitable case for this data regarding the line signal recorded
        :param lines: Line signal
        :return: LinesType enumeration
        """
        if self.use_sweeps and "Lines" in self.dict_of_data.keys():
            return LinesType.SWEEPS_REBUILD

        if self.use_sweeps and "Lines" not in self.dict_of_data.keys():
            return LinesType.SWEEPS_FROM_SCRATCH

        if "Lines" in self.dict_of_data.keys():
            lines = lines.loc[:, "abs_time"].copy()
            # Verify that the input is not corrupt
            if lines.shape[0] < self.num_of_lines // 2:
                logging.warning(
                    "Line data was corrupt - there were too few lines.\n"
                    "Simulating line data using GUI's parameters."
                )

                if "Frames" in self.dict_of_data.keys():
                    return LinesType.REBUILD
                else:
                    return LinesType.CORRUPT

            lines = self.__add_zeroth_line_event(lines=lines)

            # Analyze the line signal
            rel_idx = np.where(
                np.abs(lines.diff().pct_change(periods=1)) > self.change_thresh
            )[0]
            is_corrupt_sig: bool = (len(rel_idx) // 2) / lines.shape[
                0
            ] > self.change_thresh
            if is_corrupt_sig and "Frames" in self.dict_of_data.keys():
                return LinesType.REBUILD
            elif is_corrupt_sig and "Frames" not in self.dict_of_data.keys():
                return LinesType.CORRUPT
            else:  # not corrupt
                if self.image_soft == ImagingSoftware.SCANIMAGE.value:
                    return LinesType.SCANIMAGE
                elif self.image_soft == ImagingSoftware.MSCAN.value:
                    return LinesType.MSCAN

        else:
            return LinesType.FROM_SCRATCH

    def __si(self) -> Tuple[Dict, np.uint64]:
        return ScanImageLineValidator(sig_val=self).run()

    def __mscan(self) -> Tuple[Dict, np.uint64]:
        return MScanLineValidator(sig_val=self).run()

    def __corrupt(self) -> Tuple[Dict, int]:
        """
        Data is corrupted, and no frame channel can help us.
        :return:
        """
        self.line_delta = self.__bins_bet_lines()
        self.dict_of_data["Lines"] = self.__extrapolate_line_data(
            line_point=self.dict_of_data["Lines"].at[0, "abs_time"]
        )
        return self.dict_of_data, self.line_delta

    def __rebuild(self) -> Tuple[Dict, int]:
        """
        Data is corrupted, but we can rebuild lines on top of the frame channel
        :return:
        """
        line_array = self.__create_line_array()
        self.dict_of_data["Lines"] = pd.DataFrame(
            line_array, columns=["abs_time"], dtype="uint64"
        )
        line_delta = self.last_event_time / (
            self.num_of_lines * int(self.num_of_frames)
        )
        return self.dict_of_data, line_delta

    def __add_zeroth_line_event(self, lines: pd.Series) -> pd.Series:
        """
        If lines was used as the starting signal - add a 0 line time
        :param lines:
        :return: Corrected line signal
        """
        line_delta = lines.diff().median()
        zeroth_line_delta = np.abs(lines.iat[0] - line_delta) / line_delta
        if (
            zeroth_line_delta < 0.05
        ):  # first line came exactly line_delta time after the start of the experiment
            lines = pd.Series([0], dtype=np.uint64).append(lines)
            self.dict_of_data["Lines"] = pd.DataFrame(
                [[0] * len(self.data_to_grab)],
                columns=self.data_to_grab,
                dtype=np.uint64,
            ).append(self.dict_of_data["Lines"], ignore_index=True)

        return lines

    def __si_valid(self) -> Tuple[Dict, int]:
        """
        Data is valid. Check whether we need a 0-time line event
        """
        line_delta = self.dict_of_data["Lines"].loc[:, "abs_time"].diff().median()
        # zeroth_line_delta = np.abs(self.dict_of_data['Lines'].loc[0, 'abs_time'] - line_delta)/line_delta
        # if zeroth_line_delta < 0.05:
        #     self.dict_of_data['Lines'] = pd.DataFrame([[0] * len(self.data_to_grab)],
        #                                          columns=self.data_to_grab,
        #                                          dtype='uint64')\
        #         .append(self.dict_of_data['Lines'], ignore_index=True)
        return self.dict_of_data, np.uint64(line_delta)

    def __from_scratch(self) -> Tuple[Dict, int]:
        line_array = self.__create_line_array()
        self.dict_of_data["Lines"] = pd.DataFrame(
            line_array, columns=["abs_time"], dtype="uint64"
        )
        line_delta = self.last_event_time / (
            self.num_of_lines * int(self.num_of_frames)
        )
        return self.dict_of_data, line_delta

    def __sweeps_rebuild(self) -> Tuple[Dict, int]:
        line_delta = 1 / self.line_freq if not self.bidir else 1 / (2 * self.line_freq)
        self.num_of_frames = np.ceil(
            self.last_event_time / (line_delta / self.binwidth * self.num_of_lines)
        )
        line_array = self.__create_line_array(line_delta=line_delta / self.binwidth)
        self.dict_of_data["Lines"] = pd.DataFrame(
            line_array, columns=["abs_time"], dtype="uint64"
        )
        return self.dict_of_data, int(line_delta / self.binwidth)

    def __sweeps_from_scratch(self) -> Tuple[Dict, int]:
        sweep_vec = np.arange(self.max_sweep + 1, dtype=np.uint64)
        if len(sweep_vec) < 2:
            logging.warning(
                "All data was registered to a single sweep. Line data will be completely simulated."
            )
        else:
            self.dict_of_data["Lines"] = pd.DataFrame(
                sweep_vec * self.total_sweep_time, columns=["abs_time"], dtype=np.uint64
            )
        return self.dict_of_data, self.total_sweep_time

    def __extrapolate_line_data(self, line_point: int = 0) -> pd.DataFrame:
        """
        From a single line signal extrapolate the presumed line data vector. The line frequency is double
        the original frequency. If needed, events will be discarded later.
        :param line_point: Start interpolation from this point.
        :return: pd.DataFrame of line data
        """
        # Verify input
        if (
            line_point / self.line_delta > self.num_of_lines
        ):  # The first recorded line came very late
            line_point = 0

        # Create the matrix containing the duplicate frame data
        delay_between_frames_in_bins = int(self.delay_between_frames / self.binwidth)
        time_of_frame = (
            self.line_delta * self.num_of_lines + delay_between_frames_in_bins
        )
        num_of_frames = max(int(np.floor(self.last_event_time / time_of_frame)), 1)
        time_of_frame_mat = np.tile(
            time_of_frame * np.arange(num_of_frames), (self.num_of_lines, 1)
        )

        # Create the matrix containing the duplicate line data
        line_vec = np.arange(
            start=line_point,
            stop=self.line_delta * self.num_of_lines,
            step=self.line_delta,
            dtype=np.uint64,
        )
        line_vec = np.r_[
            np.flip(
                np.arange(
                    start=line_point, stop=0, step=-self.line_delta, dtype=np.uint64
                )[1:],
                axis=0,
            ),
            line_vec,
        ]
        # Check if 0 should be included
        if line_vec[0] - self.line_delta == 0:
            line_vec = np.r_[0, line_vec]

        line_mat = np.tile(line_vec.reshape(len(line_vec), 1), (1, num_of_frames))

        # Add them up
        assert len(line_mat) > 0
        assert len(time_of_frame_mat) > 0
        final_mat = line_mat + time_of_frame_mat
        line_vec_final = np.ravel(final_mat, order="F")

        return pd.DataFrame(line_vec_final, columns=["abs_time"], dtype=np.uint64)

    def __create_line_array(self, line_delta: int = 0) -> np.ndarray:
        """Create a pandas Series of start-of-line times"""

        total_lines = self.num_of_lines * int(self.num_of_frames)
        if line_delta == 0:
            line_array = np.arange(
                start=0,
                stop=self.last_event_time,
                step=self.last_event_time / total_lines,
                dtype=np.uint64,
            )
        else:
            line_array = np.arange(
                start=0, stop=self.last_event_time, step=line_delta, dtype=np.uint64
            )
        return line_array

    def __validate_frame_input(self):

        if "Frames" in self.dict_of_data.keys():
            if self.image_soft == "MScan":
                self.dict_of_data["Frames"] = pd.DataFrame(
                    [[0] * len(self.data_to_grab)],
                    columns=self.data_to_grab,
                    dtype="uint64",
                ).append(self.dict_of_data["Frames"], ignore_index=True)
        else:
            frame_array = self.__create_frame_array(
                lines=self.dict_of_data["Lines"].loc[:, "abs_time"]
            )
            self.dict_of_data["Frames"] = pd.DataFrame(
                frame_array, columns=["abs_time"], dtype="uint64"
            )

        return self.dict_of_data

    def __create_frame_array(self, lines: pd.Series = None) -> np.ndarray:
        """Create a pandas Series of start-of-frame times"""

        num_of_recorded_lines = lines.shape[0]
        actual_num_of_frames = max(num_of_recorded_lines // self.num_of_lines, 1)

        if num_of_recorded_lines < self.num_of_lines:
            array_of_frames = np.linspace(
                start=0,
                stop=self.last_event_time,
                num=int(actual_num_of_frames),
                endpoint=False,
                dtype=np.uint64,
            )
        else:
            unnecess_lines = num_of_recorded_lines % self.num_of_lines
            array_of_frames = lines.iloc[
                0 : int(num_of_recorded_lines - unnecess_lines) : self.num_of_lines
            ]

        return np.array(array_of_frames)

    def __validate_created_data_channels(self) -> None:
        """
        Make sure that the dictionary that contains all data channels makes sense.
        """
        assert {"PMT1", "Lines", "Frames"} <= set(
            self.dict_of_data.keys()
        )  # A is subset of B

        if (
            self.dict_of_data["Frames"].shape[0] > self.dict_of_data["Lines"].shape[0]
        ):  # more frames than lines
            raise UserWarning("More frames than lines, consider replacing the two.")

        try:
            if (
                self.dict_of_data["TAG Lens"].shape[0]
                < self.dict_of_data["Lines"].shape[0]
            ):
                raise UserWarning(
                    "More lines than TAG pulses, consider replacing the two."
                )
        except KeyError:
            pass

        try:
            if (
                self.dict_of_data["Laser"].shape[0]
                < self.dict_of_data["Lines"].shape[0]
                or self.dict_of_data["Laser"].shape[0]
                < self.dict_of_data["Frames"].shape[0]
            ):
                raise UserWarning(
                    "Laser pulses channel contained less ticks than the Lines or Frames channel."
                )
        except KeyError:
            pass

        try:
            if (
                self.dict_of_data["Laser"].shape[0]
                < self.dict_of_data["TAG Lens"].shape[0]
            ):
                raise UserWarning(
                    "Laser pulses channel contained less ticks than the TAG lens channel."
                )
        except KeyError:
            pass

    def __validate_laser_input(self, pulses) -> pd.Series:
        """
        Create an orderly laser pulse train.
        :param pulses: Laser data
        :return:
        """

        diffs = pulses.loc[:, "abs_time"].diff()
        rel_idx = (diffs <= 1.05 * np.ceil((1 / (self.laser_freq * self.binwidth)))) & (
            diffs >= 0.95 * np.floor((1 / (self.laser_freq * self.binwidth)))
        )
        pulses_final = pulses[rel_idx]  # REMINDER: Laser offset wasn't added
        if len(pulses_final) < 0.9 * len(
            pulses
        ):  # TODO: If there's holdafter time, there's a good chance that the pulses
            # will not be periodic due to this extra time.
            logging.warning(
                "More than 10% of pulses were filtered due to bad timings. Make sure the laser input is fine."
            )

        pulses_final = pd.concat(
            [pulses.loc[:0, :], pulses_final]
        )  # Add back the first pulse
        pulses_final.reset_index(drop=True, inplace=True)

        return pulses_final

    def add_phase_to_bidir_lines(self, lines: pd.Series):
        """
        "Fix" temporarily the lines for them to pass the corruption check
        :param lines: Lines after the phase change
        :param bidir_phase:
        :param binwidth:
        :return: Lines "pre-fix" of phase change - their original form
        """
        phase_in_seconds = self.bidir_phase * 1e-6
        if phase_in_seconds < 0:
            lines.iloc[1::2] -= np.uint64(np.abs(phase_in_seconds / self.binwidth))
        else:
            lines.iloc[1::2] += np.uint64(phase_in_seconds / self.binwidth)
        return lines


#!/usr/bin/env python
# encoding: utf-8

"""
A module for parsing information from various files.
"""

from __future__ import (absolute_import, division, print_function, unicode_literals)
import numpy as np
import os

from arkane.statmech import determine_qm_software
from arkane.qchem import QChemLog
from arkane.gaussian import GaussianLog
from arkane.molpro import MolproLog

from arc.common import get_logger
from arc.species.converter import get_xyz_string, standardize_xyz_string
from arc.arc_exceptions import InputError, ParserError

##################################################################

logger = get_logger()

def parse_frequencies(path, software):
    """
    Parse the frequencies from a freq job output file.
    """
    lines = _get_lines_from_file(path)
    freqs = np.array([], np.float64)
    if software.lower() == 'qchem':
        for line in lines:
            if ' Frequency:' in line:
                items = line.split()
                for i, item in enumerate(items):
                    if i:
                        freqs = np.append(freqs, [(float(item))])
    elif software.lower() == 'gaussian':
        with open(path, 'r') as f:
            line = f.readline()
            while line != '':
                if 'Frequencies --' in line:
                    freqs = np.append(freqs, [float(frq) for frq in line.split()[2:]])
                line = f.readline()
    elif software.lower() == 'molpro':
        read = False
        for line in lines:
            if 'Nr' in line and '[1/cm]' in line:
                continue
            if read:
                if line == os.linesep:
                    read = False
                    continue
                freqs = np.append(freqs, [float(line.split()[-1])])
            if 'Low' not in line and 'Vibration' in line and 'Wavenumber' in line:
                read = True
    else:
        raise ParserError('parse_frequencies() can currently only parse Molpro, QChem and Gaussian files,'
                          ' got {0}'.format(software))
    logger.debug('Using parser.parse_frequencies. Determined frequencies are: {0}'.format(freqs))
    return freqs

def parse_t1(path):
    """
    Parse the T1 parameter from a Molpro coupled cluster calculation.
    """
    lines = _get_lines_from_file(path)
    t1 = None
    for line in lines:
        if 'T1 diagnostic:' in line:
            t1 = float(line.split()[-1])
    return t1

def parse_e_elect(path, zpe_scale_factor=1.):
    """
    Parse the zero K energy, E0, from an sp job output file.
    """
    if not os.path.isfile(path):
        raise InputError('Could not find file {0}'.format(path))
    log = determine_qm_software(fullpath=path)
    try:
        e_elect = log.loadEnergy(zpe_scale_factor) * 0.001  # convert to kJ/mol
    except Exception:
        logger.warning('Could not read e_elect from {0}'.format(path))
        e_elect = None
    return e_elect

def parse_xyz_from_file(path):
    """
    Parse xyz coordinated from:
    .xyz - XYZ file
    .gjf - Gaussian input file
    .out or .log - ESS output file (Gaussian, QChem, Molpro)
    other - Molpro or QChem input file
    """
    lines = _get_lines_from_file(path)
    file_extension = os.path.splitext(path)[1]

    xyz = None
    relevant_lines = list()

    if file_extension == '.xyz':
        relevant_lines = lines[2:]
    elif file_extension == '.gjf':
        start_parsing = False
        for line in lines:
            if start_parsing and line and line != '\n' and line != '\r\n':
                relevant_lines.append(line)
            elif start_parsing:
                break
            else:
                splits = line.split()
                if len(splits) == 2 and all([s.isdigit() for s in splits]):
                    start_parsing = True
    elif 'out' in file_extension or 'log' in file_extension:
        log = determine_qm_software(fullpath=path)
        coords, number, _ = log.loadGeometry()
        xyz = get_xyz_string(coords=coords, numbers=number)
    else:
        record = False
        for line in lines:
            if '$end' in line or '}' in line:
                break
            if record and len(line.split()) == 4:
                relevant_lines.append(line)
            elif '$molecule' in line:
                record = True
            elif 'geometry={' in line:
                record = True
        if not relevant_lines:
            raise ParserError('Could not parse xyz coordinates from file {0}'.format(path))
    if xyz is None and relevant_lines:
        xyz = ''.join([line for line in relevant_lines if line])
    return standardize_xyz_string(xyz)

def parse_dipole_moment(path):
    """
    Parse the dipole moment in Debye from an opt job output file.
    """
    lines = _get_lines_from_file(path)
    log = determine_qm_software(path)
    dipole_moment = None
    if isinstance(log, GaussianLog):
        # example:
        # Dipole moment (field-independent basis, Debye):
        # X=             -0.0000    Y=             -0.0000    Z=             -1.8320  Tot=              1.8320
        read = False
        for line in lines:
            if 'dipole moment' in line.lower() and 'debye' in line.lower():
                read = True
            elif read:
                dipole_moment = float(line.split()[-1])
                read = False
    elif isinstance(log, QChemLog):
        # example:
        #     Dipole Moment (Debye)
        #          X       0.0000      Y       0.0000      Z       2.0726
        #        Tot       2.0726
        skip = False
        read = False
        for line in lines:
            if 'dipole moment' in line.lower() and 'debye' in line.lower():
                skip = True
            elif skip:
                skip = False
                read = True
            elif read:
                dipole_moment = float(line.split()[-1])
                read = False
    elif isinstance(log, MolproLog):
        # example:
        #  Dipole moment /Debye                   2.96069859     0.00000000     0.00000000
        for line in lines:
            if 'dipole moment' in line.lower() and '/debye' in line.lower():
                splits = line.split()
                dm_x, dm_y, dm_z = float(splits[-3]), float(splits[-2]), float(splits[-1])
                dipole_moment = (dm_x ** 2 + dm_y ** 2 + dm_z ** 2) ** 0.5
    else:
        raise ParserError('Currently dipole moments can only be parsed from either Gaussian, Molpro, or QChem '
                          'optimization output files')
    if dipole_moment is None:
        raise ParserError('Could not parse the dipole moment')
    return dipole_moment

def parse_polarizability(path):
    """
    Parse the polarizability from a freq job output file, returns the value in Angstrom^3.
    """
    lines = _get_lines_from_file(path)
    polarizability = None
    for line in lines:
        if 'Isotropic polarizability for W' in line:
            # example:  Isotropic polarizability for W=    0.000000       11.49 Bohr**3.
            # 1 Bohr = 0.529177 Angstrom
            polarizability = float(line.split()[-2]) * 0.529177 ** 3
    return polarizability

def _get_lines_from_file(path):
    """
    A helper function for getting a list of lines from the file at `path`.
    """
    if os.path.isfile(path):
        with open(path, 'r') as f:
            lines = f.readlines()
    else:
        raise InputError('Could not find file {0}'.format(path))
    return lines

def process_conformers_file(conformers_path):
    """
    Parse coordinates and energies from an ARC conformers file of either species or TSs.

    Args:
        conformers_path (str): The path to an ARC conformers file
                               (either a "conformers_before_optimization" or
                               a "conformers_after_optimization" file).

    Returns:
        xyz (list): Entries are optimized xyz's in a string format.
    Returns:
        energies (list): Entries float numbers representing the energies in kJ/mol.
    """
    if not os.path.isfile(conformers_path):
        raise ValueError('Conformers file {0} could not be found'.format(conformers_path))
    with open(conformers_path, 'r') as f:
        lines = f.readlines()
    xyzs, energies = list(), list()
    line_index = 0
    while line_index < len(lines):
        if 'conformer' in lines[line_index] and ':' in lines[line_index] and lines[line_index].strip()[-2].isdigit():
            xyz, energy = '', None
            line_index += 1
            print(lines[line_index].strip())
            while line_index < len(lines) and lines[line_index].strip() and 'SMILES' not in lines[line_index]\
                    and 'energy' not in lines[line_index].lower() and 'guess method' not in lines[line_index].lower():
                xyz += lines[line_index]
                line_index += 1
            while line_index < len(lines) and 'conformer' not in lines[line_index]:
                if 'relative energy:' in lines[line_index].lower():
                    energy = float(lines[line_index].split()[2])
                line_index += 1
            xyzs.append(xyz)
            energies.append(energy)
        else:
            line_index += 1
    return xyzs, energies


import tensorflow as tf
import os
import time
from datetime import datetime
from utils import *
from model import *
import numpy as np
import pdb

# ##############################################################################
# SEGMENTATION CLASS
# ##############################################################################

class SegmenterNet(object):
    def __init__(self, cfg, model_ckp_name=""):
        """ Initializes a Segmentation Model Class """

        # MODEL SETTINGS
        self.cfg = cfg

        if model_ckp_name is not "":
            self.checkpoint_file = model_ckp_name
            self.log_dir = self.cfg.log_path + "inference_results"
            if not os.path.exists(self.log_dir):
                os.makedirs(self.log_dir)
        else:
            # DIRECTORIES TO STORE OUTPUTS
            if not self.cfg.log_name == "":
                self.log_dir = self.cfg.log_path + \
                    "{}".format(datetime.now().strftime(
                        "%Y%m%d-%H%M%S")) + "_" + self.cfg.log_name
            else:
                self.log_dir = self.cfg.log_path + \
                    "{}".format(datetime.now().strftime("%Y%m%d-%H%M%S"))
            if not os.path.exists(self.log_dir):
                os.makedirs(self.log_dir)
                self.checkpoint_file = os.path.join(self.log_dir, "model.chk")
            else:
                self.checkpoint_file = os.path.join(self.log_dir, "model.chk")

        # Create log file
        log_filename = os.path.join(self.log_dir, "net_parameters.txt")
        self.log_file = open(log_filename, 'w+')

        self.init_network_model()

    def init_network_model(self):
        self.graph = tf.Graph()
        with self.graph.as_default():

            if self.cfg.log_name == 'laser':
                self.create_input_ops()
                self.logits = create_SalsaNet_laser(
                    self.input_img, self.cfg.NUM_CLASS, dropout_rate=self.cfg.DROPOUT_PROB, is_training=self.is_training)
            elif self.cfg.log_name == 'fusion_encoder':
                self.create_fusion_input_ops()
                self.logits = create_SalsaNet_encoder_fusion(
                    self.input_laser, self.input_depth, self.cfg.NUM_CLASS, dropout_rate=self.cfg.DROPOUT_PROB, is_training=self.is_training)
            elif self.cfg.log_name == 'fusion_decoder':
                self.create_fusion_input_ops()
                self.logits = create_SalsaNet_decoder_fusion(
                    self.input_laser, self.input_depth, self.cfg.NUM_CLASS, dropout_rate=self.cfg.DROPOUT_PROB, is_training=self.is_training)
            elif self.cfg.log_name == 'fusion_encoder_decoder':
                self.create_fusion_input_ops()
                self.logits = create_SalsaNet_encoder_decoder_fusion(
                    self.input_laser, self.input_depth, self.cfg.NUM_CLASS, dropout_rate=self.cfg.DROPOUT_PROB, is_training=self.is_training)
            else:
                self.create_input_ops()
                self.logits = create_SalsaNet(
                    self.input_img, self.cfg.NUM_CLASS, dropout_rate=self.cfg.DROPOUT_PROB, is_training=self.is_training)

            self.store_network_parameters()

            with tf.name_scope("preds") as scope:
                self.preds = tf.to_int32(
                    tf.argmax(self.logits, axis=-1), name=scope)

            self.create_loss_ops()
            self.create_optimization_ops()
            self.create_evaluation_metric_ops()
            self.create_summary_ops()

    def create_input_ops(self):
        with tf.variable_scope("parameters"):
            input_img_shape = (None, self.cfg.IMAGE_HEIGHT,
                               self.cfg.IMAGE_WIDTH, self.cfg.IMAGE_CHANNEL)
            output_img_shape = (None, self.cfg.IMAGE_HEIGHT,
                                self.cfg.IMAGE_WIDTH)
            self.input_img = tf.placeholder(
                tf.float32, shape=input_img_shape, name="input_img")
            self.output_img = tf.placeholder(
                tf.int32, shape=output_img_shape, name="output_img")
            self.weight_img = tf.placeholder(
                tf.float32, shape=output_img_shape, name="weight_img")
            self.dropout = tf.placeholder_with_default(
                self.cfg.DROPOUT_PROB, shape=None, name="dropout")
            self.is_training = tf.placeholder_with_default(
                False, shape=(), name="is_training")

            self.global_step = tf.Variable(
                0, name='global_step', trainable=False)
            self.learning_rate = tf.train.exponential_decay(learning_rate=self.cfg.LEARNING_RATE,
                                                            global_step=self.global_step,
                                                            decay_steps=self.cfg.LR_DECAY_CYCLE,
                                                            decay_rate=self.cfg.LR_DECAY_FACTOR,
                                                            staircase=True,
                                                            name="learningrate")

            # Create a summary to monitor the learning rate
            tf.summary.scalar("learning_rate", self.learning_rate)

    def create_fusion_input_ops(self):
        with tf.variable_scope("parameters"):
            laser_img_shape = (None, self.cfg.LASER_IMAGE_HEIGHT,
                               self.cfg.LASER_IMAGE_WIDTH, self.cfg.LASER_IMAGE_CHANNEL)
            depth_img_shape = (None, self.cfg.DEPTH_IMAGE_HEIGHT,
                               self.cfg.DEPTH_IMAGE_WIDTH, self.cfg.DEPTH_IMAGE_CHANNEL)
            output_img_shape = (None, self.cfg.DEPTH_IMAGE_HEIGHT,
                                self.cfg.DEPTH_IMAGE_WIDTH)
            self.input_laser = tf.placeholder(
                tf.float32, shape=laser_img_shape, name="input_laser")
            self.input_depth = tf.placeholder(
                tf.float32, shape=depth_img_shape, name="input_depth")
            self.output_img = tf.placeholder(
                tf.int32, shape=output_img_shape, name="output_img")
            self.weight_img = tf.placeholder(
                tf.float32, shape=output_img_shape, name="weight_img")
            self.dropout = tf.placeholder_with_default(
                self.cfg.DROPOUT_PROB, shape=None, name="dropout")
            self.is_training = tf.placeholder_with_default(
                False, shape=(), name="is_training")

            self.global_step = tf.Variable(
                0, name='global_step', trainable=False)
            self.learning_rate = tf.train.exponential_decay(learning_rate=self.cfg.LEARNING_RATE,
                                                            global_step=self.global_step,
                                                            decay_steps=self.cfg.LR_DECAY_CYCLE,
                                                            decay_rate=self.cfg.LR_DECAY_FACTOR,
                                                            staircase=True,
                                                            name="learningrate")

            # Create a summary to monitor the learning rate
            tf.summary.scalar("learning_rate", self.learning_rate)

    def store_network_parameters(self):

        self.log_file.write("\n" + ("#" * 70) + "\n" + ("#" * 29) +
                            " parameters " + ("#" * 29) + "\n" + ("#" * 70) + "\n")

        for k, v in sorted(self.cfg.items()):
            text_to_write = str(k) + " : " + str(v) + "\n"
            self.log_file.write(text_to_write)

        self.log_file.write("\n" + ("#" * 70) + "\n" + ("#" * 70) + "\n")
        self.log_file.close()

    def create_loss_ops(self):
        # LOSS FUNCTION
        with tf.variable_scope('loss') as scope:
            unrolled_logits = tf.reshape(self.logits, (-1, self.cfg.NUM_CLASS))
            unrolled_labels = tf.reshape(self.output_img, (-1,))
            unrolled_weights = tf.reshape(self.weight_img, (-1,))
            cross_entropy = tf.losses.sparse_softmax_cross_entropy(
                labels=unrolled_labels, logits=unrolled_logits, weights=unrolled_weights)
            self.loss = tf.reduce_mean(cross_entropy)

            # Create a summary to monitor the loss
            tf.summary.scalar("loss", self.loss)

    def create_optimization_ops(self):
        # OPTIMIZATION METHOD
        with tf.variable_scope('opt') as scope:
            self.optimizer = tf.train.AdamOptimizer(
                self.learning_rate, name="optimizer")
            update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
            with tf.control_dependencies(update_ops):
                self.train_op = self.optimizer.minimize(
                    self.loss, global_step=self.global_step, name="train_op")

    def create_evaluation_metric_ops(self):
        # EVALUATION METRIC - Intersection over Union IoU
        with tf.name_scope("evaluation") as scope:
            # Define the evaluation metric and update operations
            self.evaluation, self.update_evaluation_vars = tf.metrics.mean_iou(
                tf.reshape(self.output_img, [-1]),
                tf.reshape(self.preds, [-1]),
                num_classes=self.cfg.NUM_CLASS,
                name=scope)
            # Isolate metric's running variables & create their initializer and reset operators
            evaluation_vars = tf.get_collection(
                tf.GraphKeys.LOCAL_VARIABLES, scope=scope)
            self.reset_evaluation_vars = tf.variables_initializer(
                var_list=evaluation_vars)

    def create_summary_ops(self):
        with tf.name_scope('summary'):
            self.summary_writer = tf.summary.FileWriter(
                self.log_dir, graph=self.graph)
            if self.cfg.log_name == 'fusion_encoder' or self.cfg.log_name == 'fusion_decoder' or self.cfg.log_name == 'fusion_encoder_decoder':
                # only store variables belonging to the laser block
                variable_scope = tf.get_collection(
                    tf.GraphKeys.GLOBAL_VARIABLES, scope="laser_block")
                self.laser_saver = tf.train.Saver(
                    variable_scope, name="saver", max_to_keep=1)
            self.saver = tf.train.Saver(
                tf.global_variables(), name="saver", max_to_keep=1)
            self.summary_op = tf.summary.merge_all()

    def train_segmenter(self, training_data_path, validation_data_path):

        with tf.Session(graph=self.graph) as sess:
            self.initialize_vars(sess)

            for epoch in range(1, self.cfg.NUM_EPOCHS+1):
                timeStart = time.time()

                # generate batches
                training_batches, n_training_samples = generate_lidar_batch_function(
                    training_data_path, channel_nbr=self.cfg.IMAGE_CHANNEL, class_nbr=self.cfg.NUM_CLASS, loss_weights=self.cfg.CLS_LOSS_WEIGHTS, augmentation=self.cfg.DATA_AUGMENTATION)
                validation_batches, n_validation_samples = generate_lidar_batch_function(
                    validation_data_path, channel_nbr=self.cfg.IMAGE_CHANNEL, class_nbr=self.cfg.NUM_CLASS, loss_weights=self.cfg.CLS_LOSS_WEIGHTS, augmentation=self.cfg.DATA_AUGMENTATION)

                # Num batches per epoch
                n_batches = int(
                    np.ceil(n_training_samples / self.cfg.BATCH_SIZE))

                # Iterate through each mini-batch
                for step in range(n_batches):

                    # get next batch data
                    X_batch, Y_batch, W_batch = next(
                        training_batches(self.cfg.BATCH_SIZE))
                    if self.cfg.DEBUG_MODE:
                        print('X_batch', X_batch.shape, X_batch.dtype,
                              X_batch.min(), X_batch.max())
                        print('Y_batch', Y_batch.shape, Y_batch.dtype,
                              Y_batch.min(), Y_batch.max())
                        print('W_batch', W_batch.shape, W_batch.dtype,
                              W_batch.min(), W_batch.max())

                    # Runtime metadata
                    run_options = tf.RunOptions(
                        trace_level=tf.RunOptions.FULL_TRACE)
                    run_metadata = tf.RunMetadata()

                    # Train
                    feed_dict = {self.input_img: X_batch, self.output_img: Y_batch,
                                 self.weight_img: W_batch, self.is_training: True}
                    loss, _, summary = sess.run([self.loss, self.train_op, self.summary_op],
                                                feed_dict=feed_dict, options=run_options, run_metadata=run_metadata)

                    tag_name = 'epoch {} step {}'.format(epoch, step)
                    self.summary_writer.add_summary(
                        summary, n_batches*(epoch-1)+step)

                    # force tensorflow to synchronise summaries
                    self.summary_writer.flush()

                    # Print feedback every so often
                    if self.cfg.PRINT_EVERY is not None and (step+1) % self.cfg.PRINT_EVERY == 0:
                        timeElapsed = time.time() - timeStart
                        print(" EPOCH {}/{} step: {: 5d} Batch loss: {:3.5f} Time avg: {:3.5f} sec".format(
                            epoch, self.cfg.NUM_EPOCHS, step+1, loss, timeElapsed/self.cfg.PRINT_EVERY))
                        timeStart = time.time()

                # Evaluate on train and validation sets after each epoch
                train_iou, train_loss, train_ious, train_precs, train_recalls = self.evaluate(
                    training_batches, n_training_samples, sess)
                valid_iou, valid_loss, valid_ious, valid_precs, valid_recalls = self.evaluate(
                    validation_batches, n_validation_samples, sess)

                # print scores
                self.print_evaluation_scores(
                    train_iou, train_loss, train_ious, train_precs, train_recalls, tag="Training")
                self.print_evaluation_scores(
                    valid_iou, valid_loss, valid_ious, valid_precs, valid_recalls, tag="Validation")

                # keep summary data after each epoch
                self.save_summaries(sess, train_loss, train_iou, valid_loss, valid_iou, train_ious,
                                    train_precs, train_recalls, valid_ious, valid_precs, valid_recalls, epoch)

    def train_segmenter_fusion(self, training_data_path, validation_data_path):
        with tf.Session(graph=self.graph) as sess:
            self.initialize_vars(sess)

            for epoch in range(1, self.cfg.NUM_EPOCHS+1):
                timeStart = time.time()

                # generate batches
                training_batches, n_training_samples = generate_fusion_batch_function(
                    training_data_path, laser_channel_nbr=self.cfg.LASER_IMAGE_CHANNEL, depth_channel_nbr=self.cfg.DEPTH_IMAGE_CHANNEL, class_nbr=self.cfg.NUM_CLASS, loss_weights=self.cfg.CLS_LOSS_WEIGHTS, augmentation=self.cfg.DATA_AUGMENTATION)
                validation_batches, n_validation_samples = generate_fusion_batch_function(
                    validation_data_path, laser_channel_nbr=self.cfg.LASER_IMAGE_CHANNEL, depth_channel_nbr=self.cfg.DEPTH_IMAGE_CHANNEL, class_nbr=self.cfg.NUM_CLASS, loss_weights=self.cfg.CLS_LOSS_WEIGHTS, augmentation=self.cfg.DATA_AUGMENTATION)

                # Num batches per epoch
                n_batches = int(
                    np.ceil(n_training_samples / self.cfg.BATCH_SIZE))

                # Iterate through each mini-batch
                for step in range(n_batches):

                    # get next batch data
                    X1_batch, X2_batch, Y_batch, W_batch = next(
                        training_batches(self.cfg.BATCH_SIZE))

                    if self.cfg.DEBUG_MODE:
                        print('X1_batch', X1_batch.shape, X1_batch.dtype,
                              X1_batch.min(), X1_batch.max())
                        print('X2_batch', X2_batch.shape, X2_batch.dtype,
                              X2_batch.min(), X2_batch.max())
                        print('Y_batch', Y_batch.shape, Y_batch.dtype,
                              Y_batch.min(), Y_batch.max())
                        print('W_batch', W_batch.shape, W_batch.dtype,
                              W_batch.min(), W_batch.max())

                    # Runtime metadata
                    run_options = tf.RunOptions(
                        trace_level=tf.RunOptions.FULL_TRACE)
                    run_metadata = tf.RunMetadata()

                    # Train
                    feed_dict = {self.input_laser: X1_batch, self.input_depth: X2_batch, self.output_img: Y_batch,
                                 self.weight_img: W_batch, self.is_training: True}
                    loss, _, summary = sess.run([self.loss, self.train_op, self.summary_op],
                                                feed_dict=feed_dict, options=run_options, run_metadata=run_metadata)

                    tag_name = 'epoch {} step {}'.format(epoch, step)
                    self.summary_writer.add_summary(
                        summary, n_batches*(epoch-1)+step)

                    # force tensorflow to synchronise summaries
                    self.summary_writer.flush()

                    # Print feedback every so often
                    if self.cfg.PRINT_EVERY is not None and (step+1) % self.cfg.PRINT_EVERY == 0:
                        timeElapsed = time.time() - timeStart
                        print(" EPOCH {}/{} step: {: 5d} Batch loss: {:3.5f} Time avg: {:3.5f} sec".format(
                            epoch, self.cfg.NUM_EPOCHS, step+1, loss, timeElapsed/self.cfg.PRINT_EVERY))
                        timeStart = time.time()

                # Evaluate on train and validation sets after each epoch
                train_iou, train_loss, train_ious, train_precs, train_recalls = self.evaluate_fusion(
                    training_batches, n_training_samples, sess)
                valid_iou, valid_loss, valid_ious, valid_precs, valid_recalls = self.evaluate_fusion(
                    validation_batches, n_validation_samples, sess)

                # print scores
                self.print_evaluation_scores(
                    train_iou, train_loss, train_ious, train_precs, train_recalls, tag="Training")
                self.print_evaluation_scores(
                    valid_iou, valid_loss, valid_ious, valid_precs, valid_recalls, tag="Validation")

                # keep summary data after each epoch
                self.save_summaries(sess, train_loss, train_iou, valid_loss, valid_iou, train_ious,
                                    train_precs, train_recalls, valid_ious, valid_precs, valid_recalls, epoch)

    def initialize_vars(self, session):
        if tf.train.checkpoint_exists(self.checkpoint_file):
            print("- Restoring parameters from saved checkpoints")
            print("  -", self.checkpoint_file)
            self.saver.restore(session, self.checkpoint_file)
        else:
            print("Initializing weights to random values")
            session.run(tf.global_variables_initializer())
            if self.cfg.log_name == 'fusion_encoder' or self.cfg.log_name == 'fusion_decoder' or self.cfg.log_name == 'fusion_encoder_decoder':
                # add startpoint here
                laser_startpoint = '../logs/dataset_newnew/laser/model.chk-50'
                print("Initializing laser block from: ", laser_startpoint)
                self.laser_saver.restore(session, laser_startpoint)

    def predict(self, batch_data, session):
        # MAKE PREDICTIONS ON SINGLE BATCH DATA
        feed_dict = {self.input_img: batch_data, self.is_training: False}
        batch_preds = session.run(self.preds, feed_dict=feed_dict)
        preds = batch_preds.squeeze()

        return preds

    def predict_single_image(self, input_img, session):
        # MAKE PREDICTIONS ON SINGLE IMAGE DATA

        # expand image dimension
        temp_img = np.zeros(
            (1, input_img.shape[0], input_img.shape[1], input_img.shape[2]))
        temp_img[0, :, :, :] = input_img

        # MAKE PREDICTIONS ON SINGLE IMAGE
        feed_dict = {self.input_img: temp_img, self.is_training: False}
        timeStart = time.time()
        pred_img = session.run(self.preds, feed_dict=feed_dict)
        timeElapsed = (time.time() - timeStart)*1000.0
        print("predict_single_image took : {:3.5f} msec".format(timeElapsed))

        return pred_img[0]

    def predict_single_image_fusion(self, input_data, session):
        # MAKE PREDICTIONS ON SINGLE IMAGE DATA

        laser_scan = input_data[0]
        depth_img = input_data[1]

        # expand image dimension
        temp_depth = np.zeros(
            (1, depth_img.shape[0], depth_img.shape[1], 1))
        temp_depth[0, :, :, 0] = depth_img[:, :]
        temp_laser = np.zeros(
            (1, laser_scan.shape[0], laser_scan.shape[1], 1))
        temp_laser[0, :, :, 0] = laser_scan[:, :]

        # MAKE PREDICTIONS ON SINGLE IMAGE
        feed_dict = {self.input_laser: temp_laser,
                     self.input_depth: temp_depth, self.is_training: False}
        timeStart = time.time()
        pred_img = session.run(self.preds, feed_dict=feed_dict)
        timeElapsed = (time.time() - timeStart)*1000.0
        print("predict_single_image took : {:3.5f} msec".format(timeElapsed))

        return pred_img[0]

    def evaluate(self, batch_data, data_size, session):
        # EVALUATE ON BATCH DATA
        total_loss = 0
        tps = []
        fps = []
        fns = []
        n_samples = data_size
        # Num batches needed
        n_batches = int(np.ceil(n_samples/self.cfg.BATCH_SIZE))

        # Reset the running variables for evaluation metric
        session.run(self.reset_evaluation_vars)

        # Iterate through each mini-batch
        for step in range(n_batches):
            # get next batch data
            X_batch, Y_batch, W_batch = next(batch_data(self.cfg.BATCH_SIZE))

            feed_dict = {self.input_img: X_batch, self.output_img: Y_batch,
                         self.weight_img: W_batch, self.is_training: False}

            # Get loss, and update running variables for evaluation metric
            loss, preds, confusion_mtx = session.run(
                [self.loss, self.preds, self.update_evaluation_vars], feed_dict=feed_dict)
            total_loss += loss

            # iou computation
            tp, fp, fn = self.evaluate_iou(Y_batch, preds, self.cfg.NUM_CLASS)
            tps.append(tp)
            fps.append(fp)
            fns.append(fn)

        tps = np.array(tps)
        fps = np.array(fps)
        fns = np.array(fns)
        epsilon = 1e-12
        ious = tps.astype(np.float) / (tps + fns + fps + epsilon)
        precision = tps.astype(np.float) / (tps + fps + epsilon)
        recall = tps.astype(np.float) / (tps + fns + epsilon)

        mean_ious = np.mean(ious, axis=0)
        mean_prec = np.mean(precision, axis=0)
        mean_recall = np.mean(recall, axis=0)

        # Get the updated score from the running metric
        score = session.run(self.evaluation)
        # Average the loss
        avg_loss = total_loss/float(n_batches)

        return score, avg_loss, mean_ious, mean_prec, mean_recall

    def evaluate_fusion(self, batch_data, data_size, session):
        # EVALUATE ON BATCH DATA
        total_loss = 0
        tps = []
        fps = []
        fns = []
        n_samples = data_size
        # Num batches needed
        n_batches = int(np.ceil(n_samples/self.cfg.BATCH_SIZE))

        # Reset the running variables for evaluation metric
        session.run(self.reset_evaluation_vars)

        # Iterate through each mini-batch
        for step in range(n_batches):
            # get next batch data
            X1_batch, X2_batch, Y_batch, W_batch = next(
                batch_data(self.cfg.BATCH_SIZE))

            feed_dict = {self.input_laser: X1_batch, self.input_depth: X2_batch, self.output_img: Y_batch,
                         self.weight_img: W_batch, self.is_training: False}

            # Get loss, and update running variables for evaluation metric
            loss, preds, confusion_mtx = session.run(
                [self.loss, self.preds, self.update_evaluation_vars], feed_dict=feed_dict)
            total_loss += loss

            # iou computation
            tp, fp, fn = self.evaluate_iou(Y_batch, preds, self.cfg.NUM_CLASS)
            tps.append(tp)
            fps.append(fp)
            fns.append(fn)

        tps = np.array(tps)
        fps = np.array(fps)
        fns = np.array(fns)
        epsilon = 1e-12
        ious = tps.astype(np.float) / (tps + fns + fps + epsilon)
        precision = tps.astype(np.float) / (tps + fps + epsilon)
        recall = tps.astype(np.float) / (tps + fns + epsilon)

        mean_ious = np.mean(ious, axis=0)
        mean_prec = np.mean(precision, axis=0)
        mean_recall = np.mean(recall, axis=0)

        # Get the updated score from the running metric
        score = session.run(self.evaluation)
        # Average the loss
        avg_loss = total_loss/float(n_batches)

        return score, avg_loss, mean_ious, mean_prec, mean_recall

    def evaluate_iou(self, label, pred, n_class):

        assert label.shape == pred.shape, \
            'label and pred shape mismatch: {} vs {}'.format(
                label.shape, pred.shape)

        tps = np.zeros(n_class)
        fns = np.zeros(n_class)
        fps = np.zeros(n_class)

        for cls_id in range(n_class):
            tp = np.sum(pred[label == cls_id] == cls_id)
            fp = np.sum(label[pred == cls_id] != cls_id)
            fn = np.sum(pred[label == cls_id] != cls_id)

            tps[cls_id] = tp
            fps[cls_id] = fp
            fns[cls_id] = fn

        return tps, fps, fns

    def expand_image_dimension(self, input_img):
        # return pred results as colored rgb images
        n_samples = input_img.shape[0]
        output_img = np.zeros(
            (n_samples, input_img.shape[1], input_img.shape[2], 3))
        for i in range(0, n_samples):
            label_map = input_img[i, :, :]
            color_img = np.zeros((input_img.shape[1], input_img.shape[2], 3))
            for j in range(0, self.cfg.NUM_CLASS):
                color_img[label_map == j, :] = self.cfg.CLS_COLOR_MAP[j]

            output_img[i, :, :, :] = color_img

        return output_img

    def save_summaries(self, sess, train_loss, train_iou, valid_loss, valid_iou, train_mean_ious, train_precs, train_recalls, valid_mean_ious, valid_precs, valid_recalls, epoch):

        # Save checkpoints
        self.saver.save(sess, self.checkpoint_file,
                        global_step=epoch, write_meta_graph=True)

        # Save training and validation summaries
        summary = tf.Summary()
        summary.value.add(tag='Training/Training Loss',
                          simple_value=float(train_loss))
        summary.value.add(tag='Validation/Validation Loss',
                          simple_value=float(valid_loss))
        summary.value.add(tag='Training/Training IOU',
                          simple_value=float(train_iou))
        summary.value.add(tag='Validation/Validation IOU',
                          simple_value=float(valid_iou))

        for i in range(0, self.cfg.NUM_CLASS):
            tag_name = 'Training/' + self.cfg.CLASSES[i] + '/IOU'
            summary.value.add(
                tag=tag_name, simple_value=float(train_mean_ious[i]))
            tag_name = 'Training/' + self.cfg.CLASSES[i] + '/Prec'
            summary.value.add(tag=tag_name, simple_value=float(train_precs[i]))
            tag_name = 'Training/' + self.cfg.CLASSES[i] + '/Recall'
            summary.value.add(
                tag=tag_name, simple_value=float(train_recalls[i]))

            tag_name = 'Validation/' + self.cfg.CLASSES[i] + '/IOU'
            summary.value.add(
                tag=tag_name, simple_value=float(valid_mean_ious[i]))
            tag_name = 'Validation/' + self.cfg.CLASSES[i] + '/Prec'
            summary.value.add(tag=tag_name, simple_value=float(valid_precs[i]))
            tag_name = 'Validation/' + self.cfg.CLASSES[i] + '/Recall'
            summary.value.add(
                tag=tag_name, simple_value=float(valid_recalls[i]))

        self.summary_writer.add_summary(summary, epoch)

        # force tensorflow to synchronise summaries
        self.summary_writer.flush()

    def print_evaluation_scores(self, iou, loss, ious, precs, recalls, tag="Training"):

        if tag == "Training":
            s = "TR IOU: {: 3.3f} TR IOU: {: 3.3f} TR LOSS: {: 3.5f} "
        elif tag == "Testing":
            s = "TEST IOU: {: 3.3f} TEST IOU: {: 3.3f} TEST LOSS: {: 3.5f} "
        else:
            s = "VR IOU: {: 3.3f} VR IOU: {: 3.3f} VR LOSS: {: 3.5f} "
        print(s.format(iou, np.mean(ious), loss))

        for i in range(0, self.cfg.NUM_CLASS):
            s = self.cfg.CLASSES[i] + " PREC: {: 3.3f} " + self.cfg.CLASSES[i] + \
                " REC: {: 3.3f} " + self.cfg.CLASSES[i] + " IOU: {: 3.3f}"
            print(s.format(precs[i], recalls[i], ious[i]))


from dataclasses import dataclass
from typing import List, Tuple

import numpy as np
import pkg_resources

@dataclass
class Solo8BaseConfig:
  dt: float = 1e-3
  # Max torque supplied by the motors
  motor_torque_limit: float = 2

  # TODO: Figure out how to lay the robot flat so that it doesn't need to fall
  robot_start_pos: Tuple[float] = (0., 0., 0.5)
  robot_start_orientation_euler: Tuple[float] = (0., 0., 0.)
  gravity: Tuple[float] = (0., 0., -9.81)

  max_motor_rotation = 2 * np.pi

  linear_damping: float = .04
  angular_damping: float = .04
  restitution: float = 0.
  lateral_friction: float = 0.5

  render_width: int = 369
  render_height: int = 369
  render_fov: int = 80
  render_aspect: float = render_width / render_height
  render_pos = [0, 0, .2]
  render_cam_distance = 1
  render_yaw = 0.
  render_pitch = -20.
  render_roll = 0.

  @property
  def urdf(self):
    return pkg_resources.resource_filename('gym_solo', self.urdf_path)


#===================================
#        Chengwen Liu              #
#      liuchw2010@gmail.com        #
#   University of Texas at Austin  #
#===================================

import os
import sys
import subprocess
import numpy as np

def main():
  usage = ''' Usage: python ldensity.py xxx.arc '''
  arcfile = sys.argv[1]
  if len(sys.argv) != 2:
    sys.exit(usage)
  if not os.path.isfile(arcfile):
    sys.exit(f"{arcfile} does not exist!")

  massdict = {"H": 1.007941, "C":12.01074, "N":14.006703, "O":15.999405, "S":32.0648, "P":30.973761998, \
              "F": 18.99840, "CL":35.4529, "BR":79.9035, "NA":22.989769, "K":39.0983, "MG":24.3051, "ZN":65.38}

  natom = subprocess.check_output(["head", "-n1", arcfile]).split()[0]
  headstr = "head -n%s %s > tmp.xyz "%(int(natom)+2, arcfile)
  os.system(headstr)
  grepstr = "grep ' 90.000000' %s > lattice.dat"%(arcfile)
  os.system(grepstr)
  atoms = np.loadtxt("tmp.xyz", usecols=(1,), unpack=True, dtype='str', skiprows=2)
  a,b,c = np.loadtxt("lattice.dat", usecols=(0,1,2), unpack=True, dtype='float')
  os.system("rm -rf tmp.xyz lattice.dat")
  mass = 0.0
  navogadro = 6.02214086 #x10**23
  for atom in atoms:
    mass+=massdict[atom[:1].upper()]
  sumdens = 0.0
  volume = (a*b*c).mean()
  dens = 1000*10.0*(mass/navogadro)/volume
  print("%20s%10.1f kg/m^3" %("Average Density:", dens))
  return

if __name__ == "__main__":
  main()


# Copyright 2019 Graphcore Ltd.
import numpy as np
import torch
from torch import nn
import pytest

import popart
import onnx

from bert_model import BertConfig, Bert, ExecutionMode, get_model
from bert import set_library_seeds
from tests.torch_bert import BertConfig as TorchBertConfig, BertOutput, BertIntermediate
from tests.utils import run_py, copy_weights_to_torch, run_fwd_model, check_tensors, check_model

'''
Tests the fully connected layers.
'''
test_modes = [ExecutionMode.DEFAULT, ExecutionMode.PINGPONG]

def simplified_gelu(x):
    return x * torch.sigmoid(1.702 * x)

class BertFCN(nn.Module):
    def __init__(self, config):
        super(BertFCN, self).__init__()
        self.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)

    def forward(self, input_x):
        intermediate_output = self.intermediate(input_x)
        layer_output = self.output(intermediate_output, input_x)
        # add attentions if we output them
        outputs = (layer_output, )
        return outputs

ACTIVATIONS = {
    'Relu': ['Relu', nn.functional.relu],
    'Gelu': ['Gelu', nn.functional.gelu],
    'SGelu': ['SGelu', simplified_gelu],
    'Leaky_relu':
    ['LRelu', lambda x: nn.functional.leaky_relu(x, negative_slope=0.1)]
}

TORCH_TO_ONNX = {
    ExecutionMode.DEFAULT: {
        "intermediate.dense.weight": "1/W",
        "intermediate.dense.bias": "1/B",
        "output.dense.weight": "2/W",
        "output.dense.bias": "2/B",
        "output.LayerNorm.weight": "Gamma",
        "output.LayerNorm.bias": "Beta"
    },
    ExecutionMode.PINGPONG: {
        "intermediate.dense.weight": "FF/1/Dense/Weight",
        "intermediate.dense.bias": "FF/1/Dense/Bias",
        "output.dense.weight": "FF/2/Dense/Weight",
        "output.dense.bias": "FF/2/Dense/Bias",
        "output.LayerNorm.weight": "FF/Norm/Gamma",
        "output.LayerNorm.bias": "FF/Norm/Beta"
    }
}

TRANSPOSE_WEIGHTS = {
    "intermediate.dense.weight": np.transpose,
    "output.dense.weight": np.transpose
}

@pytest.mark.parametrize('mode', test_modes)
@pytest.mark.parametrize('activation_function', ACTIVATIONS.keys())
@pytest.mark.parametrize('phase', ('fwd', 'bwd'))
def test_activation_function(mode, activation_function, phase):

    set_library_seeds(0)

    popart_act_function, pytorch_activation = ACTIVATIONS[activation_function]
    config = BertConfig(vocab_length=128,
                        batch_size=1,
                        hidden_size=768,
                        sequence_length=128,
                        popart_dtype="FLOAT",
                        no_dropout=True,
                        activation_type=str(popart_act_function))

    data, outputs, proto, post_proto = popart_result_and_model(
        config, mode, is_bwd=False if phase is 'fwd' else True)

    inputs = [
        data.reshape(config.batch_size, config.sequence_length,
                     config.hidden_size)
    ]

    #  ------------------- PyTorch -------------------------
    torch_config = TorchBertConfig(config.vocab_length,
                                   config.hidden_size,
                                   config.num_layers,
                                   config.attention_heads,
                                   layer_norm_eps=config.layer_norm_eps,
                                   hidden_dropout_prob=0.,
                                   hidden_act=pytorch_activation)

    torch_output, torch_model = pytorch_result_and_model(
        torch_config,
        inputs,
        proto,
        mode,
        is_bwd=False if phase is 'fwd' else True)

    check_tensors(torch_output, outputs, margin=1e-07)

    if phase is 'bwd':
        check_model(torch_model,
                    post_proto,
                    TORCH_TO_ONNX[mode],
                    transform=TRANSPOSE_WEIGHTS,
                    margin=1e-07)

def popart_result_and_model(popart_config, mode, is_bwd=False):
    popart_model = get_model(popart_config, mode, 'feedforward')

    input_info = popart.TensorInfo(popart_config.popart_dtype, [
        popart_config.batch_size * popart_config.sequence_length,
        popart_config.hidden_size
    ])
    input_tensor = popart_model.builder.addInputTensor(input_info)

    data = {
        input_tensor:
        np.random.normal(0, 0.02,
                         input_info.shape()).astype(popart_config.dtype)
    }

    user_options = {}
    if mode == ExecutionMode.PINGPONG:
        user_options = {
            "batchSerializationFactor": 1,
            "pingPongPhases": popart_model.total_ping_pong_phases
        }
        output = popart_model(input_tensor)
    else:
        user_options = {"enableStochasticRounding": True}
        output = popart_model.feed_forward(input_tensor)

    if is_bwd:
        l1_lambda = 0.1
        if mode == ExecutionMode.PINGPONG:
            with popart_model.scope_provider(popart_model.builder, popart_model.norm.scope):
                l1 = popart_model.builder.aiGraphcore.l1loss(
                    [output],
                    l1_lambda,
                    debugPrefix="l1LossVal",
                    reduction=popart.ReductionType.Sum)

        else:
            l1 = popart_model.builder.aiGraphcore.l1loss(
                [output],
                l1_lambda,
                debugPrefix="l1LossVal",
                reduction=popart.ReductionType.Sum)
        proto = popart_model.builder.getModelProto()
        optimizer = popart.ConstSGD(0.01)

        outputs, post_proto = run_py(proto,
                                     data, (output, l1),
                                     loss=l1,
                                     optimizer=optimizer,
                                     user_options=user_options,
                                     execution_mode=mode)
    else:
        proto = popart_model.builder.getModelProto()
        outputs, post_proto = run_py(proto,
                                     data,
                                     output,
                                     user_options=user_options,
                                     execution_mode=mode)

    return data[input_tensor], outputs, proto, post_proto

def pytorch_result_and_model(torch_config,
                             inputs,
                             popart_proto,
                             mode,
                             is_bwd=False):
    # Conversion of the popart model to onnx
    proto = onnx.load_model_from_string(popart_proto)

    torch_model = BertFCN(torch_config)
    # Turn off dropout
    torch_model.eval()

    copy_weights_to_torch(torch_model,
                          proto,
                          TORCH_TO_ONNX[mode],
                          transform=TRANSPOSE_WEIGHTS)

    result = run_fwd_model(inputs, torch_model)

    if is_bwd:
        l1_lambda = 0.1
        optim = torch.optim.SGD(torch_model.parameters(),
                                0.01,
                                weight_decay=0.0,
                                momentum=0.0)

        result = torch_model(*[torch.from_numpy(t).float() for t in inputs])[0]
        torch_loss = l1_lambda * torch.norm(result, 1)
        torch_loss.backward()
        optim.step()
        result = result.detach().numpy()

    return result, torch_model