"""FDTD Postprocessing.

"""
from __future__ import print_function
from builtins import str
from builtins import range
from builtins import object

import os
import numpy
import EMpy.utils
import pylab

__author__ = "Lorenzo Bolla"

class Input(object):
    """Data structure to handle input files."""

    def __init__(self, filename):
        """Set the input filename."""
        self.filename = filename

    def __str__(self):
        """Return a representation of the input file."""

        dftmon_str = "%g ! #timemonitors \n" % len(self.dftmonitors)
        if len(self.dftmonitors) > 0:
            dftmon_str += "".join(
                [
                    "%g %g %g %g %g %g\n%g %g\n"
                    % (
                        dm[0][0],
                        dm[0][1],
                        dm[0][2],
                        dm[0][3],
                        dm[0][4],
                        dm[0][5],
                        dm[1][0],
                        dm[1][1],
                    )
                    for dm in self.dftmonitors
                ]
            )

        timemon_str = "%g ! #timemonitors \n" % len(self.dftmonitors)
        if len(timemon_str) > 0:
            timemon_str += "%g %g \n %s" % (
                self.timemonitors_time_interval[0],
                self.timemonitors_time_interval[1],
                "".join(
                    [
                        "%g %g %g ! time_monitor #%d\n" % (s[0], s[1], s[2], iss)
                        for iss, s in enumerate(self.timemonitors)
                    ]
                ),
            )

        return (
            "%g %g %g %g ! dx dy dz cfl \n"
            "%g %g %g %g %g %g %s %g %g ! xmax ymax zmax pmlx pmly pmlz pmltype pmlsmooth pmlref \n"
            "%g %g %g %g ! xmax ymax zmax pmlx pmly pmlz \n"
            "%g ! output3deps? \n"
            "%g ! number diel slices \n"
            "%s \n"
            "%g ! number field slices \n"
            "%s \n"
            "%g %g %g ! #dielobjs, index of bg, conductivity of bg \n"
            "%s"
            "%g ! smoothing method \n"
            "%g ! #sources \n"
            "%s"
            "%g %g %g ! lambdamin, lambdamax, dlambda \n"
            "%s"
            "%s"
            % (
                self.dx,
                self.dy,
                self.dz,
                self.cfl,
                self.xmax,
                self.ymax,
                self.zmax,
                self.pmlx,
                self.pmly,
                self.pmlz,
                self.pmltype,
                self.pmlsmooth,
                self.pmlref,
                self.start,
                self.end,
                self.slides,
                self.snapshot,
                self.output3deps,
                len(self.dielslices),
                "\n".join(
                    [
                        "%g %g %g ! dielslice #%d" % (d[0], d[1], d[2], dd)
                        for (dd, d) in enumerate(self.dielslices)
                    ]
                ),
                len(self.fieldslices),
                "\n".join(
                    [
                        "%g %g %g ! fieldslice #%d" % (f[0], f[1], f[2], ff)
                        for (ff, f) in enumerate(self.fieldslices)
                    ]
                ),
                len(self.dielobjs),
                self.bgrix,
                self.bgsigma,
                "".join(["%s %s\n" % obj for obj in self.dielobjs]),
                self.smoothing_method,
                len(self.sources),
                "".join(["%s\n%s\n%s\n%s\n" % src for src in self.sources]),
                self.lambdamin,
                self.lambdamax,
                self.dlambda,
                dftmon_str,
                timemon_str,
            )
        )

    def tofile(self, filename=None):
        """Save the input data to the input file."""
        if filename is None:
            filename = self.filename
        f = open(filename, "w")
        f.write(self.__str__())
        f.close()

class Param(object):
    """Data structure to handle the param file."""

    def __str__(self):
        """Return a representation of the input file."""
        return (
            "%g\n"
            "%g\n"
            "%g\n"
            "%g\n"
            "%g\n"
            "%g\n"
            "%g\n"
            "%g\n"
            "%g\n"
            "%g\n"
            "%g\n"
            "%g\n"
            "%g\n"
            "%g\n"
            "%s"
            % (
                self.dx,
                self.dy,
                self.dz,
                self.dt,
                self.mx,
                self.my,
                self.mz,
                self.pmlx,
                self.pmly,
                self.pmlz,
                self.nflux,
                self.ntime,
                self.step1,
                self.step2,
                "\n".join(
                    [
                        "%d\n%d\n%d\n%d\n%d\n%d"
                        % (
                            dm["direction"],
                            dm["nfreq"],
                            dm["flxlim"][0],
                            dm["flxlim"][1],
                            dm["flxlim"][2],
                            dm["flxlim"][3],
                        )
                        for dm in self.dftmonitors
                    ]
                ),
            )
        )

class Sensor(object):
    """Data structure to handle the FFT sensor's data."""

    def plot(self, n):
        """Plot the sensor's fields."""
        pylab.clf()
        pylab.hot()
        pylab.subplot(2, 2, 1)
        pylab.contour(numpy.abs(self.E1[:, :, n].T), 16)
        pylab.axis("image")
        pylab.title("E1")
        pylab.subplot(2, 2, 2)
        pylab.contour(numpy.abs(self.H1[:, :, n].T), 16)
        pylab.axis("image")
        pylab.title("H1")
        pylab.subplot(2, 2, 3)
        pylab.contour(numpy.abs(self.E2[:, :, n].T), 16)
        pylab.axis("image")
        pylab.title("E2")
        pylab.subplot(2, 2, 4)
        pylab.contour(numpy.abs(self.H2[:, :, n].T), 16)
        pylab.axis("image")
        pylab.title("H2")
        pylab.show()

    def __str__(self):
        """Return a representation of the sensor."""
        return "E1\n%s\nH1\n%s\nE2\n%s\nH2\n%s\n" % (self.E1, self.H1, self.E2, self.H2)

class TimeSensor(object):
    """Data structure to handle the time sensor's data."""

    def plot_Ex(self, logplot=False):
        self.__plot_field(self.Ex, logplot)

    def plot_Ey(self, logplot=False):
        self.__plot_field(self.Ey, logplot)

    def plot_Ez(self, logplot=False):
        self.__plot_field(self.Ez, logplot)

    def plot_Hx(self, logplot=False):
        self.__plot_field(self.Hx, logplot)

    def plot_Hy(self, logplot=False):
        self.__plot_field(self.Hy, logplot)

    def plot_Hz(self, logplot=False):
        self.__plot_field(self.Hz, logplot)

    def __plot_field(self, field, logplot=False):
        if logplot:
            data = 20 * numpy.log10(1e-20 + numpy.abs(field))
            pylab.plot(self.t, data)
        else:
            data = field
            pylab.plot(self.t, data)
        pylab.show()

class FDTD(object):
    """FDTD.
    Data structure to handle an FDTD simulation. It manages an input file, a param file and the sensors' output.
    It can run a simulation via a system call.
    """

    def __init__(self):
        self.input = None
        self.param = None
        self.sensors = None

    def fetch_data(
        self,
        remote_dir_="./",
        input_file="inp.txt",
        param_file="param",
        directory_="./",
    ):
        remote_dir = fixdir(remote_dir_)
        directory = fixdir(directory_)
        # input file
        os.system(
            "scp -C bollalo001@pico:" + remote_dir + "/" + input_file + " " + directory
        )
        # param file
        os.system(
            "scp -C bollalo001@pico:" + remote_dir + "/" + param_file + " " + directory
        )
        # fieldslices, flux and time sensors
        os.system(
            "scp -C bollalo001@pico:" + remote_dir + "/[EHeh]*_*" + " " + directory
        )
        # dielslices
        os.system("scp -C bollalo001@pico:" + remote_dir + "/diel*" + " " + directory)

    def put_data(self, remote_dir_="./", input_file="inp.txt", directory_="./"):
        remote_dir = fixdir(remote_dir_)
        directory = fixdir(directory_)
        # input file
        os.system("scp -C" + directory + input_file + " bollalo001@pico:" + remote_dir)
        # .dat modesolver's files
        os.system("scp -C" + directory + "*.dat bollalo001@pico:" + remote_dir)

    def load(
        self, directory_="./", input_file="inp.txt", param_file="param", remote_dir_=""
    ):
        """Load input, param and sensors."""
        remote_dir = fixdir(remote_dir_)
        directory = fixdir(directory_)
        if remote_dir != "":
            self.fetch_data(remote_dir, input_file, param_file, directory)
        self.load_input_file(directory, input_file)
        self.load_param(directory, param_file)
        self.load_sensors(directory)
        self.load_time_sensors(directory)

    def load_input_file(self, directory_="./", filename="inp.txt"):
        """Load input file."""
        directory = fixdir(directory_)
        try:
            f = open(directory + filename)
        except Exception:
            print("ERROR: input file")
            return
        inp = Input(filename)

        (inp.dx, inp.dy, inp.dz, inp.cfl) = numpy.fromstring(
            strip_comment(f.readline()), sep=" "
        )
        tmp = strip_comment(f.readline())
        tmp_idx = tmp.find("P")
        if tmp_idx > 0:
            inp.pmltype = "P"
        else:
            tmp_idx = tmp.find("G")
            if tmp_idx > 0:
                inp.pmltype = "G"
            else:
                raise ValueError("wrong pmltype")
        (inp.xmax, inp.ymax, inp.zmax, inp.pmlx, inp.pmly, inp.pmlz) = numpy.fromstring(
            tmp[:tmp_idx], sep=" "
        )
        (inp.pmlsmooth, inp.pmlref) = numpy.fromstring(tmp[tmp_idx + 1 :], sep=" ")
        (inp.start, inp.end, inp.slides, inp.snapshot) = numpy.fromstring(
            strip_comment(f.readline()), sep=" "
        )
        inp.output3deps = numpy.fromstring(strip_comment(f.readline()), sep=" ")

        # dielslices
        ndielslices = numpy.fromstring(strip_comment(f.readline()), sep=" ")
        inp.dielslices = []
        for i in range(ndielslices):
            inp.dielslices.append(
                numpy.fromstring(strip_comment(f.readline()), sep=" ")
            )

        # fieldslices
        nfieldslices = numpy.fromstring(strip_comment(f.readline()), sep=" ")
        inp.fieldslices = []
        for i in range(nfieldslices):
            inp.fieldslices.append(
                numpy.fromstring(strip_comment(f.readline()), sep=" ")
            )

        # dielobjs
        (ndielobjs, inp.bgrix, inp.bgsigma) = numpy.fromstring(
            strip_comment(f.readline()), sep=" "
        )
        inp.dielobjs = []
        for i in range(int(ndielobjs)):
            inp.dielobjs.append(
                (strip_comment(f.readline()), strip_comment(f.readline()))
            )
        inp.smoothing_method = numpy.fromstring(strip_comment(f.readline()), sep=" ")

        # sources
        nsources = numpy.fromstring(strip_comment(f.readline()), dtype=int, sep=" ")
        inp.sources = []
        #        (inp.time_dependence, inp.wls, inp.pwidth, inp.shift) = numpy.fromstring(strip_comment(f.readline()), sep = ' ')
        for i in range(nsources):
            inp.sources.append(
                (
                    strip_comment(f.readline()),
                    strip_comment(f.readline()),
                    strip_comment(f.readline()),
                    strip_comment(f.readline()),
                )
            )

        # dft monitors
        (inp.lambdamin, inp.lambdamax, inp.dlambda) = numpy.fromstring(
            strip_comment(f.readline()), sep=" "
        )
        ndftmonitors = numpy.fromstring(strip_comment(f.readline()), dtype=int, sep=" ")
        inp.dftmonitors = []
        for i in range(ndftmonitors):
            inp.dftmonitors.append(
                (
                    numpy.fromstring(strip_comment(f.readline()), sep=" "),
                    numpy.fromstring(strip_comment(f.readline()), sep=" "),
                )
            )

        # time monitors
        ntimemonitors = numpy.fromstring(strip_comment(f.readline()), sep=" ")
        inp.timemonitors_time_interval = numpy.fromstring(
            strip_comment(f.readline()), sep=" "
        )
        inp.timemonitors = []
        for i in range(ntimemonitors):
            inp.timemonitors.append(
                numpy.fromstring(strip_comment(f.readline()), sep=" ")
            )

        f.close()
        self.input = inp

    def load_param(self, directory_="./", filename="param"):
        """Load param file."""
        directory = fixdir(directory_)
        param = Param()
        try:
            data = numpy.fromfile(directory + filename, sep=" ")
        except Exception:
            print("ERROR: param file")
            return
        param.dx, param.dy, param.dz, param.dt = data[0:4]
        (
            param.mx,
            param.my,
            param.mz,
            param.pmlx,
            param.pmly,
            param.pmlz,
            param.nflux,
            param.ntime,
            param.step1,
            param.step2,
        ) = data[4:14].astype(numpy.int32)
        param.dftmonitors = []
        for iflux in range(int(param.nflux)):
            direction, nfreq = data[14 + iflux * 6 : 16 + iflux * 6]
            flxlim = data[16 + iflux * 6 : 20 + iflux * 6]
            param.dftmonitors.append(
                {"direction": int(direction), "nfreq": int(nfreq), "flxlim": flxlim}
            )
        self.param = param

    def load_time_sensors(self, directory_="./"):
        """Load time sensors."""
        directory = fixdir(directory_)
        time_sensors = []
        if self.param is None:
            self.load_param(directory)
        for itime in range(self.param.ntime):
            tmp = TimeSensor()
            tmp.Ex = load_fortran_unformatted(directory + "Ex_time_%02d" % (itime + 1))
            tmp.Ey = load_fortran_unformatted(directory + "Ey_time_%02d" % (itime + 1))
            tmp.Ez = load_fortran_unformatted(directory + "Ez_time_%02d" % (itime + 1))
            tmp.Hx = load_fortran_unformatted(directory + "Hx_time_%02d" % (itime + 1))
            tmp.Hy = load_fortran_unformatted(directory + "Hy_time_%02d" % (itime + 1))
            tmp.Hz = load_fortran_unformatted(directory + "Hz_time_%02d" % (itime + 1))
            tmp.t = self.param.dt * numpy.arange(len(tmp.Ex))
            time_sensors.append(tmp)

        self.time_sensors = time_sensors

    def load_sensors(self, directory_="./"):
        """Load sensors."""
        directory = fixdir(directory_)
        sensors = []
        if self.param is None:
            self.load_param(directory)
        for iflux in range(self.param.nflux):
            tmp = Sensor()
            dm = self.param.dftmonitors[iflux]
            tmp.E1 = load_fortran_unformatted(directory + "E1_%02d" % (iflux + 1))
            tmp.H1 = load_fortran_unformatted(directory + "H1_%02d" % (iflux + 1))
            tmp.E2 = load_fortran_unformatted(directory + "E2_%02d" % (iflux + 1))
            tmp.H2 = load_fortran_unformatted(directory + "H2_%02d" % (iflux + 1))
            # [tmp.E1, tmp.H1, tmp.E2, tmp.H2] = map(lambda x: x[0::2] + 1j * x[1::2], [tmp.E1, tmp.H1, tmp.E2, tmp.H2])
            # more memory efficient!
            tmp.E1 = tmp.E1[0::2] + 1j * tmp.E1[1::2]
            tmp.H1 = tmp.H1[0::2] + 1j * tmp.H1[1::2]
            tmp.E2 = tmp.E2[0::2] + 1j * tmp.E2[1::2]
            tmp.H2 = tmp.H2[0::2] + 1j * tmp.H2[1::2]

            n1 = dm["flxlim"][1] - dm["flxlim"][0] + 1
            n2 = dm["flxlim"][3] - dm["flxlim"][2] + 1
            tmp.E1 = tmp.E1.reshape((n1, n2 + 1, dm["nfreq"]), order="F")
            tmp.H1 = tmp.H1.reshape((n1, n2 + 1, dm["nfreq"]), order="F")
            tmp.E2 = tmp.E2.reshape((n1 + 1, n2, dm["nfreq"]), order="F")
            tmp.H2 = tmp.H2.reshape((n1 + 1, n2, dm["nfreq"]), order="F")
            if dm["direction"] == 1:
                # sensors in the x-direction
                tmp.dx1 = self.param.dy
                tmp.dx2 = self.param.dz
            elif dm["direction"] == 2:
                # sensors in the y-direction
                tmp.dx1 = self.param.dx
                tmp.dx2 = self.param.dz
            elif dm["direction"] == 3:
                # sensors in the z-direction
                tmp.dx1 = self.param.dx
                tmp.dx2 = self.param.dy
            else:
                raise ValueError("wrong direction")

            sensors.append(tmp)

        self.sensors = sensors

    def viz2D(self, filename, directory_="./", const_dir="z", logplot=False):
        """Visualize a slice."""
        directory = fixdir(directory_)
        data = load_fortran_unformatted(directory + filename)
        if self.param is None:
            self.load_param(directory)
        x = numpy.linspace(
            self.param.dx / 2.0,
            self.param.dx * self.param.mx - self.param.dx / 2.0,
            self.param.mx,
        )
        y = numpy.linspace(
            self.param.dy / 2.0,
            self.param.dy * self.param.my - self.param.dy / 2.0,
            self.param.my,
        )
        z = numpy.linspace(
            self.param.dz / 2.0,
            self.param.dz * self.param.mz - self.param.dz / 2.0,
            self.param.mz,
        )
        if const_dir == "x":
            n1 = self.param.my
            n2 = self.param.mz
            x1 = y
            x2 = z
            x1label = "y"
            x2label = "z"
        elif const_dir == "y":
            n1 = self.param.mx
            n2 = self.param.mz
            x1 = x
            x2 = z
            x1label = "x"
            x2label = "z"
        else:
            n1 = self.param.mx
            n2 = self.param.my
            x1 = x
            x2 = y
            x1label = "x"
            x2label = "y"
        data = data.reshape((n2, n1))
        pylab.clf()
        if logplot:
            data = 20 * numpy.log10(numpy.abs(data).clip(1e-30, 1e30))
            pylab.jet()
        else:
            pylab.hot()
        pylab.contour(x1, x2, data, 64)
        pylab.colorbar()
        pylab.axis("image")
        pylab.xlabel(x1label + " /um")
        pylab.ylabel(x2label + " /um")
        pylab.show()

    def memory(self):
        """Estimate the memory occupation."""
        # size_of_char = 1
        # size_of_int = 4
        size_of_real = 4
        # size_of_complex = 2 * size_of_real
        # size_of_dielobj = size_of_int + 31 * size_of_real + 2 * 16 * size_of_char
        # size_of_source = 9 * size_of_int + 5 * size_of_real + 6 * 16 * size_of_char
        # size_of_monitor = (6 + 2) * 6 * size_of_int

        Gb = 1024 ** 3
        max_available_RAM = 32 * Gb

        dynamic_alloc_memory = 0

        # epsx, epsy, epsz
        dynamic_alloc_memory = (
            dynamic_alloc_memory
            + 3
            * (self.param.mx + 2 * self.input.pmlx)
            * (self.param.my + 2 * self.input.pmly)
            * (self.param.mz + 2 * self.input.pmlz)
            * size_of_real
        )
        # sigma
        dynamic_alloc_memory = (
            dynamic_alloc_memory
            + 2
            * (self.param.mx + 2 * self.input.pmlx)
            * (self.param.my + 2 * self.input.pmly)
            * (self.param.mz + 2 * self.input.pmlz)
            * size_of_real
        )
        # cex, cmx
        dynamic_alloc_memory = (
            dynamic_alloc_memory
            + 2 * 2 * (self.param.mx + 2 * self.input.pmlx) * size_of_real
        )
        # cey, cmy
        dynamic_alloc_memory = (
            dynamic_alloc_memory
            + 2 * 2 * (self.param.my + 2 * self.input.pmly) * size_of_real
        )
        # cez, cmz
        dynamic_alloc_memory = (
            dynamic_alloc_memory
            + 2 * 2 * (self.param.mz + 2 * self.input.pmlz) * size_of_real
        )

        # exy, exz, eyx, eyz, ...
        dynamic_alloc_memory = (
            dynamic_alloc_memory
            + 12
            * (self.param.mx + 2 * self.input.pmlx)
            * (self.param.my + 2 * self.input.pmly)
            * (self.param.mz + 2 * self.input.pmlz)
            * size_of_real
        )

        print(
            "Alloc mem = %g Gb, [%d%%]"
            % (
                1.0 * dynamic_alloc_memory / Gb,
                int(1.0 * dynamic_alloc_memory / max_available_RAM * 100),
            )
        )

    def run(
        self,
        directory_="./",
        exe_file="/xlv1/labsoi_devices/devices/f3d",
        output_file="output",
        ncpu=12,
        bg=False,
        remote=True,
    ):
        """Run the simulation, possibly in remote."""
        directory = fixdir(directory_)
        #        os.environ['OMP_NUM_THREAD'] = str(ncpu)
        #        cmd = 'dplace -x6 ' + exe_file + ' > ' + output_file
        cmd = (
            "cd"
            + directory
            + "; setenv OMP_NUM_THREAD"
            + str(ncpu)
            + "dplace -x6 "
            + exe_file
            + " > "
            + output_file
        )
        if bg:
            cmd += "&"
        if remote:
            cmd = 'ssh pico "' + cmd + '"'
        os.system(cmd)

    def __str__(self):
        """Return a representation of the FDTD data structure."""
        return "INPUT\n%s\nPARAM\n%s\nSENSORS\n%s\n" % (
            self.input,
            self.param,
            self.sensors,
        )

def load_fortran_unformatted(filename):
    """Load data from an unformatted fortran binary file."""
    try:
        f = open(filename, "rb")
    except Exception:
        print("ERROR")
        return
    nbytes = numpy.fromfile(f, dtype=numpy.int32, count=1)
    n = nbytes / numpy.float32().nbytes
    data = numpy.fromfile(f, dtype=numpy.float32, count=n)
    f.close()
    return data

def strip_comment(line):
    """Get rid of fortran comments."""
    idx = line.find("!")
    if idx != -1:
        return line[:idx].strip()
    return line

def fixdir(str, sep="/"):
    tmp = str
    if len(tmp) > 0:
        if tmp[-1] != sep:
            tmp += sep
    return tmp

# def overlap_f(simul, solver, nwl):
#     vu = numpy.zeros((len(simul.sensors), len(solver.modes)), dtype=complex)
#     ju = numpy.zeros((len(simul.sensors), len(solver.modes)), dtype=complex)
#     for isens, sens in enumerate(simul.sensors):
#         for imode, mode in enumerate(solver.modes):
#             Ex, Ey, Ez, Hx, Hy, Hz = mode.get_fields_for_FDTD()
#             vu[isens, imode] = 0.5 * (
#                         numpy.trapz(numpy.trapz(sens.E1[:,1:-1,nwl] * Hy, dx=sens.dx2*1e-6), dx=sens.dx1*1e-6) -
#                         numpy.trapz(numpy.trapz(sens.E2[1:-1,:,nwl] * Hx, dx=sens.dx2*1e-6), dx=sens.dx1*1e-6))
#             ju[isens, imode] = 0.5 * (
#                         numpy.trapz(numpy.trapz(sens.H2[1:-1,:,nwl] * Ey, dx=sens.dx2*1e-6), dx=sens.dx1*1e-6) -
#                         numpy.trapz(numpy.trapz(sens.H1[:,1:-1,nwl] * Ex, dx=sens.dx2*1e-6), dx=sens.dx1*1e-6))
#     A = (vu + ju) / 2.
#     B = (vu - ju) / 2.
#     Pm = numpy.abs(A)**2 - numpy.abs(B)**2
#     P = Pm.sum(axis=1)
#     return (vu, ju, A, B, Pm, P)

def overlap_f(sensors, solver, nwl):
    vu = numpy.zeros((len(sensors), len(solver.modes)), dtype=complex)
    ju = numpy.zeros((len(sensors), len(solver.modes)), dtype=complex)
    for isens, sens in enumerate(sensors):
        x = sens.dx1 * numpy.arange(sens.E2.shape[0])
        y = sens.dx2 * numpy.arange(sens.E1.shape[1])
        for imode, mode in enumerate(solver.modes):
            # resample the mode to the sensor's grid
            Ex, Ey, Ez, Hx, Hy, Hz = mode.get_fields_for_FDTD(x, y)
            #            vu[isens, imode] = 0.5 * (
            #                        numpy.trapz(numpy.trapz(sens.E1[:,1:-1,nwl] * Hy, dx=sens.dx2*1e-6), dx=sens.dx1*1e-6) -
            #                        numpy.trapz(numpy.trapz(sens.E2[1:-1,:,nwl] * Hx, dx=sens.dx2*1e-6), dx=sens.dx1*1e-6))
            #            ju[isens, imode] = 0.5 * (
            #                        numpy.trapz(numpy.trapz(sens.H2[1:-1,:,nwl] * Ey, dx=sens.dx2*1e-6), dx=sens.dx1*1e-6) -
            #                        numpy.trapz(numpy.trapz(sens.H1[:,1:-1,nwl] * Ex, dx=sens.dx2*1e-6), dx=sens.dx1*1e-6))
            vu[isens, imode] = 0.5 * (
                EMpy.utils.trapz2(
                    sens.E1[:, 1:-1, nwl] * Hy, dx=sens.dx1 * 1e-6, dy=sens.dx2 * 1e-6
                )
                - EMpy.utils.trapz2(
                    sens.E2[1:-1, :, nwl] * Hx, dx=sens.dx1 * 1e-6, dy=sens.dx2 * 1e-6
                )
            )
            ju[isens, imode] = 0.5 * (
                EMpy.utils.trapz2(
                    sens.H2[1:-1, :, nwl] * Ey, dx=sens.dx1 * 1e-6, dy=sens.dx1 * 1e-6
                )
                - EMpy.utils.trapz2(
                    sens.H1[:, 1:-1, nwl] * Ex, dx=sens.dx1 * 1e-6, dy=sens.dx1 * 1e-6
                )
            )
    A = (vu + ju) / 2.0
    B = (vu - ju) / 2.0
    Pm = numpy.abs(A) ** 2 - numpy.abs(B) ** 2
    P = Pm.sum(axis=1)
    return (vu, ju, A, B, Pm, P)


#Arguments in order:
#    1.) hoomd_path (string): the path to hoomd on your machine
#    2.) tsteps (int): the number of timesteps your simulation is running
#    3.) dump_freq (int): the frequency of gsd dumps
#    4.) part_frac_a (float): particle fraction of type a particles
#    5.) pe_a (float): activity of type a particles
#    6.) pe_b (float): activity of tyep b particles

import sys
import os

#hoomd_path = str(sys.argv[1])
hoomd_path = "${hoomd_path}"
#tsteps = int(sys.argv[2])
runfor = ${runfor}
#dump_freq = int(sys.argv[3])
dump_freq = ${dump_freq}
#part_frac_a = float(sys.argv[4])
part_perc_a = ${part_frac_a}
part_frac_a = float(part_perc_a) / 100.0
#pe_a = float(sys.argv[5])
pe_a = ${pe_a}
#pe_b = float(sys.argv[6])
pe_b = ${pe_b}
#gsd_path = str(sys.argv[7])
gsd_path = "${gsd_path}"
#part_num = int(sys.argv[8])
part_num = ${part_num}
#phi = float(sys.argv[9])
phi = ${phi}
phi = float(phi)/100.0
#seed1 = int(sys.argv[10])
#seed2 = int(sys.argv[11])
#seed3 = int(sys.argv[12])
#seed4 = int(sys.argv[13])
#seed5 = int(sys.argv[14])
seed1 = ${seed1}                # seed for position
seed2 = ${seed2}                # seed for bd equilibration
seed3 = ${seed3}                # seed for initial orientations
seed4 = ${seed4}                # seed for A activity
seed5 = ${seed5}                # seed for B activity

# tau = sigma^2 / diffusion coefficient
tau = 1
# dt = 2E-5 * tau, or, x * sigma^2 / Diffusion coeff
my_dt = 0.000001 * tau
# run for 100 tau, 100 * sigma^2 / Diffusion coeff
sim_length = runfor * tau
# compute number of tsteps to achieve this
tsteps = sim_length / my_dt
#my_dt = 0.00005
#tsteps = 50000000
# calculate number of tsteps which are dumped
dumps = tsteps/dump_freq

sys.path.append(hoomd_path)

import hoomd
from hoomd import md
from hoomd import deprecated
import numpy as np

pow = np.log10(1/my_dt)
one_length = int(18*(pow-2)+29)             # gives length of array w/ values below 1 tau
tau_to_tstep = tau / my_dt                  # this is 1 tau in terms of tsteps
spacer = tau / (10*my_dt)                   # 1/10th of tau, the spacer
gr_one_len = (tsteps - tau_to_tstep)/spacer # gives length of remaining array (tau > 1)
ar_tot_len = int(gr_one_len + one_length)

#get tsteps for msd calculations, needs to be in tau
msd_dumps = np.zeros((ar_tot_len), dtype=np.float64)
jumper = 5
value_to_dump = 15
count = 10
for iii in range(0,len(msd_dumps)):
    if iii <= 10:
        msd_dumps[iii] = iii
    elif value_to_dump * my_dt >= 1:
        msd_dumps[iii] = tau_to_tstep
        tau_to_tstep += spacer
    elif count == 95:
        msd_dumps[iii] = value_to_dump
        jumper *= 10
        value_to_dump += jumper
        count = 10
    else:
        msd_dumps[iii] = value_to_dump
        value_to_dump += jumper
        count += 5

ten_size = 0
for jjj in range(0,len(msd_dumps)):
    if msd_dumps[jjj]*my_dt <= 10:
        ten_size += 1

msd_ten = np.zeros((ten_size), dtype=np.float64)
for hhh in range(0,len(msd_ten)):
    msd_ten[hhh] = msd_dumps[hhh]

msd_dumps += 110000
last_ten = 10 * tau / my_dt
msd_ten += tsteps - last_ten + 110000

#initialize system randomly, can specify GPU execution here

hoomd.context.initialize()
# attempting to initialize a high density with min_dist > 0.75 is unsuccesful
system = hoomd.deprecated.init.create_random(N = part_num,
                                             phi_p = phi,
                                             name = 'A',
                                             min_dist = 0.7,
                                             seed = seed1,
                                             dimensions = 2)

system.particles.types.add('B')
snapshot = system.take_snapshot()

part_a = part_num * part_frac_a         # get the total number of A particles
part_a = int(part_a)
part_b = part_num - part_a              # get the total number of B particles
part_b = int(part_b)
mid = int(part_a)                       # starting point for assigning B particles

if part_perc_a == 0:                    # take care of all b case
    mid = 0
    for i in range(mid,part_num):
        system.particles[i].type = 'B'
elif part_perc_a != 100:                # mix of each
    for i in range(mid,part_num):
        system.particles[i].type = 'B'

all=hoomd.group.all()
gA = hoomd.group.type(type = 'A', update=True)
gB = hoomd.group.type(type = 'B', update=True)
N = len(all)
Na = len(gA)
Nb = len(gB)

#define potential between pairs
nl = hoomd.md.nlist.cell()
lj = hoomd.md.pair.lj(r_cut=2**(1/6), nlist=nl)
lj.set_params(mode='shift')
lj.pair_coeff.set('A', 'A', epsilon=1.0, sigma=1.0)
lj.pair_coeff.set('A', 'B', epsilon=1.0, sigma=1.0)
lj.pair_coeff.set('B', 'B', epsilon=1.0, sigma=1.0)

#integrator type
#hoomd.md.integrate.mode_minimize_fire(group=all, dt=0.00001, ftol=1e-2, Etol=1e-7)
#fire = hoomd.md.integrate.mode_minimize_fire(dt=0.00001, ftol=1e-2, Etol=1e-7)
#nve = integrate.nve(group=all)
#hoomd.run(1000)
#nve.disable()

#run simulation with current settings here
#hoomd.md.integrate.mode_standard(dt=0.000005)
#hoomd.md.integrate.mode_standard(dt=0.0000005)
hoomd.md.integrate.mode_standard(dt=my_dt)

# Here's where you'll have to differentiate between particle types
# -each group should get a different temperature
# -Pe = 3 v_0 tau_r / sigma  OR  = v_0 sigma / D_t
# -tau_r = 1 / D_r
# -D_r = 3 D_t / sigma^2
# -D_t = sigma^2 / tau_lj
# -D_t = kbT / 3*pi*eta*sigma
# -tau_lj = sigma^2 / epsilon * beta * D_t
# -tau_brown = sigma^2 / D_t
hoomd.md.integrate.brownian(group=all, kT=1.0, seed=seed2)
hoomd.run(100000)

#set the activity of each type
np.random.seed(seed3)                           # seed for random orientations
angle = np.random.rand(part_num) * 2 * np.pi    # random number for particle orientation

if part_perc_a != 0 and part_perc_a != 100:
    activity_a = []
    for i in range(0,mid):
        x = (np.cos(angle[i])) * pe_a
        y = (np.sin(angle[i])) * pe_a
        z = 0
        tuple = (x, y, z)
        activity_a.append(tuple)
    activity_b = []
    for i in range(mid,part_num):
        x = (np.cos(angle[i])) * pe_b
        y = (np.sin(angle[i])) * pe_b
        z = 0
        tuple = (x, y, z)
        activity_b.append(tuple)
    hoomd.md.force.active(group=gA,
                          seed=seed4,
                          f_lst=activity_a,
                          rotation_diff=3.0,
                          orientation_link=False,
                          orientation_reverse_link=True)
    hoomd.md.force.active(group=gB,
                          seed=seed5,
                          f_lst=activity_b,
                          rotation_diff=3.0,
                          orientation_link=False,
                          orientation_reverse_link=True)
else:
    if part_perc_a == 0:
        activity_b = []
        for i in range(0,part_num):
            x = (np.cos(angle[i])) * pe_b
            y = (np.sin(angle[i])) * pe_b
            z = 0
            tuple = (x, y, z)
            activity_b.append(tuple)
        hoomd.md.force.active(group=gB,
                              seed=seed5,
                              f_lst=activity_b,
                              rotation_diff=3.0,
                              orientation_link=False,
                              orientation_reverse_link=True)
    else:
        activity_a = []
        for i in range(0,part_num):
            x = (np.cos(angle[i])) * pe_a
            y = (np.sin(angle[i])) * pe_a
            z = 0
            tuple = (x, y, z)
            activity_a.append(tuple)
        hoomd.md.force.active(group=gA,
                              seed=seed4,
                              f_lst=activity_a,
                              rotation_diff=3.0,
                              orientation_link=False,
                              orientation_reverse_link=True)

#write dumps
name = "pa" + str(pe_a) + "_pb" + str(pe_b) + "_xa" + str(part_perc_a) + ".gsd"
msd_name = "MSD_pa" + str(pe_a) + "_pb" + str(pe_b) + "_xa" + str(part_perc_a) + ".gsd"
msd_tentau = "MSDten_pa" + str(pe_a) + "_pb" + str(pe_b) + "_xa" + str(part_perc_a) + ".gsd"
sqlite_name = "pa" + str(pe_a) + "_pb" + str(pe_b) + "_xa" + str(part_perc_a) + ".sqlite"

### Dump for MSD ###
# change this to units of tau?
#def dump_spec(timestep):
#
#    if timestep in msd_dumps:
#        hoomd.dump.gsd(filename=msd_name,
#                       period=None,
#                       group=all,
#                       overwrite=False,
#                       dynamic=['attribute', 'property', 'momentum'])
#        os.close(2)
#
#    if timestep in msd_ten:
#        hoomd.dump.gsd(filename=msd_tentau,
#                       period=None,
#                       group=all,
#                       overwrite=False,
#                       dynamic=['attribute', 'property', 'momentum'])
#        os.close(2)
#
#hoomd.analyze.callback(callback = dump_spec, period = 1)
####################

hoomd.dump.gsd(name,
               period=dump_freq,
               group=all,
               overwrite=True,
               phase=-1,
               dynamic=['attribute', 'property', 'momentum'])

hoomd.dump.getar.simple(sqlite_name, dump_freq, 'a',
                        static=['dimensions', 'viz_static'],
                        dynamic=['viz_aniso_dynamic', 'virial', 'velocity'])

# Don't have to instantiate the compute, already passed to integrator
#hoomd.compute.thermo(group=gB)
#hoomd.analyze.log(filename="pressure_outii.txt", quantities=["pressure"], period=dump_freq)

#run
hoomd.run(tsteps)

#########################################################################
########################## Begin Data Analysis ##########################
#########################################################################

sys.path.append(gsd_path)
import gsd
from gsd import hoomd
from gsd import pygsd
import numpy as np

myfile = "pa" + str(pe_a) + "_pb" + str(pe_b) + "_xa" + str(part_perc_a) + ".gsd"
msdfile = "MSD_pa" + str(pe_a) + "_pb" + str(pe_b) + "_xa" + str(part_perc_a) + ".gsd"

f = hoomd.open(name=myfile, mode='rb')
dumps = f.__len__()
size_min = 1000                                         # minimum size of cluster

position_array = np.zeros((dumps), dtype=np.ndarray)    # array of position arrays
type_array = np.zeros((dumps), dtype=np.ndarray)        # particle types
box_data = np.zeros((1), dtype=np.ndarray)              # box dimensions
timesteps = np.zeros((dumps), dtype=np.float64)         # timesteps

with hoomd.open(name=myfile, mode='rb') as t:           # open for reading
    snap = t[0]                                         # snap 0th snapshot
    box_data = snap.configuration.box                   # get box dimensions
    for i in range(0,dumps):
        snap = t[i]                                     # take snap of each dump
        type_array[i] = snap.particles.typeid
        position_array[i] = snap.particles.position     # store all particle positions
        timesteps[i] = snap.configuration.step          # store tstep for plotting purposes

timesteps -= timesteps[0]
msd_time = timesteps[1:]

pos_A = np.zeros((dumps), dtype=np.ndarray)             # type A positions
pos_B = np.zeros((dumps), dtype=np.ndarray)             # type B positions
tmpA = np.zeros((part_a, 3), dtype=np.float32)          # temporary storage arrays
tmpB = np.zeros((part_b, 3), dtype=np.float32)

from freud import parallel, box, density, cluster
parallel.setNumThreads(1)                               # don't run multiple threads

my_density = density.LocalDensity(r_cut=2.5,
                                  volume=0.79,
                                  diameter=1.0)         # initiate class, use area of circle

l_box = box_data[0]                                     # get box dimensions (square here)
f_box = box.Box(Lx=l_box,
                Ly=l_box,
                is2D=True)                               # initialize freud box

my_clusters = cluster.Cluster(box=f_box,
                              rcut=1.0)                 # initialize class
cluster_props = cluster.ClusterProperties(box=f_box)

number_clusters = np.zeros((dumps), dtype=np.ndarray)   # arrays to store things
ids = np.zeros((dumps), dtype=np.ndarray)
size_clusters = np.zeros((dumps), dtype=np.ndarray)
tot_size = np.zeros((dumps), dtype=np.ndarray)          # number of particles in clusters
tot_num = np.zeros((dumps), dtype=np.ndarray)           # total number of clusters
MCS = np.zeros((dumps), dtype=np.ndarray)               # Mean cluster size
GF = np.zeros((dumps), dtype=np.ndarray)                # Gas fraction
A_ids = np.zeros((part_a), dtype=np.ndarray)            # type A ids
B_ids = np.zeros((part_b), dtype=np.ndarray)            # type B ids
percent_A = np.zeros((dumps), dtype=np.ndarray)         # composition A at each timestep
largest = np.zeros((dumps), dtype=np.ndarray)           # read out largest cluster at each tstep

LIQ_A = np.zeros((dumps - 1), dtype=np.ndarray)         # arrays for MSD
LIQ_B = np.zeros((dumps - 1), dtype=np.ndarray)
GAS_A = np.zeros((dumps - 1), dtype=np.ndarray)
GAS_B = np.zeros((dumps - 1), dtype=np.ndarray)
MSD_T = np.zeros((dumps - 1), dtype=np.float64)
MSD_TL = np.zeros((dumps - 1), dtype=np.ndarray)
MSD_TG = np.zeros((dumps - 1), dtype=np.ndarray)

disp_x = np.zeros((part_num), dtype=np.ndarray)         # displacement vectors
disp_y = np.zeros((part_num), dtype=np.ndarray)
disp_z = np.zeros((part_num), dtype=np.ndarray)

# analyze all particles
for j in range(0, dumps):

    l_pos = position_array[j]
    my_clusters.computeClusters(l_pos)
    number_clusters[j] = my_clusters.getNumClusters()   # find number of clusters
    ids = my_clusters.getClusterIdx()                   # get cluster ids
    cluster_props.computeProperties(l_pos, ids)
    size_clusters[j] = cluster_props.getClusterSizes()  # get number of particles in each

    how_many = my_clusters.getNumClusters()

    #############################################################
    ### This finds the cluster ids for type A and B particles ###
    #############################################################
#    A_id_count = 0
#    B_id_count = 0
#    for h in range(0, part_num):
#        if type_array[j][h] == 0:
#            A_ids[A_id_count] = ids[h]                  # store the cluster ids for A type
#            A_id_count += 1                             # IMPROVE: sort while placing?
#        else:
#            B_ids[B_id_count] = ids[h]                  # store the cluster ids for B type
#            B_id_count += 1                             # could put ids in order ...
#
#    clust_dat = np.zeros((how_many), dtype = np.ndarray)
#    clust_dat_A = np.zeros((how_many), dtype = np.ndarray)
#    clust_dat_B = np.zeros((how_many), dtype = np.ndarray)
#    numerator_A = 0
#    denominator_tot = 0

    #######################################################################
    ### If clusters are greater than a threshold size, find composition ###
    #######################################################################

#    for m in range(0, how_many):
#        clust_dat_A[m] = (A_ids == m).sum()             # sum all A type particles in a cluster
#        clust_dat_B[m] = (B_ids == m).sum()
#        clust_dat[m] = clust_dat_A[m] + clust_dat_B[m]  # find total number of particles in cluster
#        if clust_dat[m] > 15:
#            numerator_A += clust_dat_A[m]
#            denominator_tot += clust_dat[m]
#    # get the total percent of A particles in all clusters
#    if denominator_tot != 0:
#        percent_A[j] =  float(numerator_A) / float(denominator_tot)

    #####################################################################
    ### Find avg cluster size, gas fraction, and largest cluster size ###
    #####################################################################
    l_clust = 0                                             # int size of largest cluster
    for k in range(0, len(size_clusters[j])):
        # the size minimum is a very important value to consider
        if size_clusters[j][k] > size_min and size_clusters[j][k] < part_num:
            tot_size[j] += size_clusters[j][k]
            tot_num[j] += 1
            if size_clusters[j][k] > l_clust:           # if larger cluster is found
                l_clust = size_clusters[j][k]           # set l_clust to that size

    largest[j] = l_clust                                # save largest cluster size for tstep

    f_largest = "largest_pa" + str(pe_a) + "_pb" + str(pe_b) + "_xa" + str(part_perc_a) + ".txt"
    if j == 0:
        a_w = 'w'
    else:
        a_w = 'a'
    f = open(f_largest, a_w)
    f.write(str(l_clust) + '\n')
    f.close()

    if tot_num[j] > 0:
        MCS[j] = float(tot_size[j]/tot_num[j])/float(part_num)
        GF[j] = float(part_num - tot_size[j]) / float(part_num)
    else:
        MCS[j] = 0
        GF[j] = 1

    #########################################################
    ### Find MSD for A, B individually, also total system ###
    #########################################################

    # ?you can enhance difference between gas and liq by setting min clust requirement

    sort_id = np.sort(ids)                              # array of IDs sorted small to large
    q_clust = np.zeros((how_many), dtype=np.ndarray)    # my binary 'is it clustered?' array
    index = 0                                           # index of the sorted array to look at
    for a in range(0,len(q_clust)):
        add_clust = 0
        while 1:
            add_clust += 1
            if index == part_num:                       # break if index is too large
                break
            if sort_id[index] != a:                     # break if ID changes
                break
            if add_clust == 1:                          # all particles appear once
                q_clust[a] = 0
            if add_clust > size_min:                    # only multiple ids appear twice
                q_clust[a] = 1
            index += 1                                  # increment index

    lq_a_count = 0
    lq_b_count = 0
    gs_a_count = 0
    gs_b_count = 0
    if j > 0:
        numerator_A = 0
        denominator_tot = 0
        for b in range(0,part_num):

            # check instantaneous disp. over last timestep
            dx = position_array[j][b][0] - position_array[j-1][b][0]
            dy = position_array[j][b][1] - position_array[j-1][b][1]
            dz = position_array[j][b][2] - position_array[j-1][b][2]

            # if it is over some threshold, then it went past a boundary
            if dx < -50:
                dx += l_box
            if dx > 50:
                dx -= l_box
            disp_x[b] += dx

            if dy < -50:
                dy += l_box
            if dy > 50:
                dy -= l_box
            disp_y[b] += dy

            if dz < -50:
                dz += l_box
            if dz > 50:
                dz -= l_box
            disp_z[b] += dz

            msd_val = np.sqrt(((disp_x[b])**2) + ((disp_y[b])**2) + ((disp_z[b])**2))
            MSD_T[j-1] += msd_val
            if q_clust[ids[b]] == 1:                        # check if in liquid
                MSD_TL[j-1] += msd_val                      # add to tot. lq. msd
                if type_array[j][b] == 0:                   # type A case
                    LIQ_A[j-1] += msd_val
                    lq_a_count += 1
                else:
                    LIQ_B[j-1] += msd_val
                    lq_b_count += 1
            else:                                           # else, particle is gas
                MSD_TG[j-1] += msd_val                      # add to tot. gs. msd
                if type_array[j][b] == 0:                   # type A case
                    GAS_A[j-1] += msd_val
                    gs_a_count += 1
                else:
                    GAS_B[j-1] += msd_val
                    gs_b_count += 1

        # if-gating these so we don't break our program
        if lq_a_count != 0: LIQ_A[j-1] /= lq_a_count
        if lq_b_count != 0: LIQ_B[j-1] /= lq_b_count
        if gs_a_count != 0: GAS_A[j-1] /= gs_a_count
        if gs_b_count != 0: GAS_B[j-1] /= gs_b_count
        MSD_T[j-1] /= part_num
        if lq_a_count + lq_b_count != 0: MSD_TL[j-1] /= lq_a_count + lq_b_count
        if gs_a_count + gs_b_count != 0: MSD_TG[j-1] /= gs_a_count + gs_b_count

        numerator_A = lq_a_count
        denominator_tot = lq_a_count + lq_b_count

        if denominator_tot != 0:
            percent_A[j] =  float(numerator_A) / float(denominator_tot)

############################
### Density caluclations ###
############################

def getDensityPlease(n):                                # call this function as needed
    l_pos = position_array[n]                           # get ith position array
    my_density.compute(f_box,
                       l_pos,
                       l_pos)
    return my_density.getDensity()

avg_sys_density = np.zeros((1), dtype=np.ndarray)

take_last = dumps - 10
last = dumps - 1
msd_last = dumps - 2
for j in range(take_last, dumps):
    avg_sys_density[0] += getDensityPlease(j)

avg_sys_density[0] /= (dumps - take_last)

################################################################################
###### perform the same analysis on species A and species B individually #######
################################################################################

if part_perc_a != 0 and part_perc_a != 100:

    tot_size_A = np.zeros((dumps), dtype=np.ndarray)          # number of particles in clusters
    tot_num_A = np.zeros((dumps), dtype=np.ndarray)           # total number of clusters
    MCS_A = np.zeros((dumps), dtype=np.ndarray)               # Mean cluster size
    GF_A = np.zeros((dumps), dtype=np.ndarray)                # Gas fraction

    tot_size_B = np.zeros((dumps), dtype=np.ndarray)          # number of particles in clusters
    tot_num_B = np.zeros((dumps), dtype=np.ndarray)           # total number of clusters
    MCS_B = np.zeros((dumps), dtype=np.ndarray)               # Mean cluster size
    GF_B = np.zeros((dumps), dtype=np.ndarray)                # Gas fraction

    for j in range(0, dumps):
        countA = 0
        countB = 0
        for g in range(0, part_num):
            if type_array[j][g] == 0:
                tmpA[countA][0] = position_array[j][g][0]
                tmpA[countA][1] = position_array[j][g][1]
                tmpA[countA][2] = position_array[j][g][2]
                countA += 1
            else:
                tmpB[countB][0] = position_array[j][g][0]
                tmpB[countB][1] = position_array[j][g][1]
                tmpB[countB][2] = position_array[j][g][2]
                countB += 1

        pos_A[j] = tmpA
        pos_B[j] = tmpB

        l_pos = pos_A[j]
        my_clusters.computeClusters(l_pos)
        number_clusters[j] = my_clusters.getNumClusters()   # find number of clusters
        ids = my_clusters.getClusterIdx()                   # get cluster ids
        cluster_props.computeProperties(l_pos, ids)
        size_clusters[j] = cluster_props.getClusterSizes()  # get number of particles in each

        ####################################
        ### GF, MCS for A-A correlations ###
        ####################################

        for k in range(0, len(size_clusters[j])):
            # the size minimum is a very important value to consider
            if size_clusters[j][k] > size_min and size_clusters[j][k] < part_num:
                tot_size_A[j] += size_clusters[j][k]
                tot_num_A[j] += 1

        if tot_num_A[j] > 0:
            MCS_A[j] = float(tot_size_A[j]/tot_num_A[j])/float(part_a)
            GF_A[j] = float(part_a - tot_size_A[j]) / float(part_a)

        else:
            MCS_A[j] = 0
            GF_A[j] = 1

        l_pos = pos_B[j]
        my_clusters.computeClusters(l_pos)
        number_clusters[j] = my_clusters.getNumClusters()   # find number of clusters
        ids = my_clusters.getClusterIdx()                   # get cluster ids
        cluster_props.computeProperties(l_pos, ids)
        size_clusters[j] = cluster_props.getClusterSizes()  # get number of particles in each

        ####################################
        ### GF, MCS for A-A correlations ###
        ####################################

        for k in range(0, len(size_clusters[j])):
            # the size minimum is a very important value to consider
            if size_clusters[j][k] > size_min and size_clusters[j][k] < part_num:
                tot_size_B[j] += size_clusters[j][k]
                tot_num_B[j] += 1

        if tot_num_B[j] > 0:
            MCS_B[j] = float(tot_size_B[j]/tot_num_B[j])/float(part_b)
            GF_B[j] = float(part_b - tot_size_B[j]) / float(part_b)

        else:
            MCS_B[j] = 0
            GF_B[j] = 1

    def getDensityA(n):                                     # call this function as needed
        countA = 0
        for g in range(0, part_num):
            if type_array[n][g] == 0:
                tmpA[countA][0] = position_array[n][g][0]
                tmpA[countA][1] = position_array[n][g][1]
                tmpA[countA][2] = position_array[n][g][2]
                countA += 1
        pos_A[n] = tmpA
        l_pos = pos_A[n]                                    # get ith position array
        my_density.compute(f_box,
                           l_pos,
                           l_pos)
        return my_density.getDensity()

    avg_dense_A = np.zeros((1), dtype=np.ndarray)

    for j in range(take_last, dumps):
        avg_dense_A[0] += getDensityA(j)

    avg_dense_A[0] /= (dumps - take_last)

    def getDensityB(n):                                     # call this function as needed
        countB = 0
        for g in range(0, part_num):
            if type_array[n][g] == 1:
                tmpB[countB][0] = position_array[n][g][0]
                tmpB[countB][1] = position_array[n][g][1]
                tmpB[countB][2] = position_array[n][g][2]
                countB += 1
        pos_B[n] = tmpB
        l_pos = pos_B[n]                                    # get ith position array
        my_density.compute(f_box,
                           l_pos,
                           l_pos)
        return my_density.getDensity()

    avg_dense_B = np.zeros((1), dtype=np.ndarray)

    for j in range(take_last, dumps):
        avg_dense_B[0] += getDensityB(j)

    avg_dense_B[0] /= (dumps - take_last)

################################################################################
#################### Plot the individual and total data ########################
################################################################################

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns
sns.set(color_codes=True)

plt_name  = "pa" + str(pe_a) + "_pb" + str(pe_b) + "_xa" + str(part_perc_a)
plt_name1 = "pa" + str(pe_a) + "_pb" + str(pe_b) + "_xa" + str(part_perc_a) + "A"
plt_name2 = "pa" + str(pe_a) + "_pb" + str(pe_b) + "_xa" + str(part_perc_a) + "B"

if part_perc_a != 0 and part_perc_a != 100:
    sns.kdeplot(avg_sys_density[0], shade = True, color="g")
    sns.kdeplot(avg_dense_A[0], shade = True, color="r")
    sns.kdeplot(avg_dense_B[0], shade = True, color="b")
    plt.savefig('avg_density_' + plt_name + '.png', dpi=1000)
    plt.close()

    sns.kdeplot(getDensityPlease(last), shade = True, color="g")
    sns.kdeplot(getDensityA(last), shade = True, color="r")
    sns.kdeplot(getDensityB(last), shade = True, color="b")
    plt.savefig('final_density_' + plt_name + '.png', dpi=1000)
    plt.close()

    plt.plot(MCS, color="g")
    plt.plot(MCS_A, color="r")
    plt.plot(MCS_B, color="b")
    #plt.ylim((0,1))
    plt.savefig('MCS_'+ plt_name + '.png', dpi=1000)
    plt.close()

    plt.plot(GF, color="g")
    plt.plot(GF_A, color="r")
    plt.plot(GF_B, color="b")
    plt.ylim((0,1))
    plt.savefig('GF_'+plt_name+'.png', dpi=1000)
    plt.close()

    plt.plot(percent_A, color="r")
    #plt.ylim((0,1))
    plt.savefig('A_comp_'+plt_name+'.png', dpi=1000)
    plt.close()

    plt.plot(largest, color="g")
    plt.savefig('Largest_clust_'+plt_name+'.png', dpi=1000)
    plt.close()

    plt.plot(msd_time, GAS_A,  color="r", marker='o', markersize=1, linestyle='None', label='Gas_A')
    plt.plot(msd_time, GAS_B,  color="b", marker='o', markersize=1, linestyle='None', label='Gas_B')
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('Timesteps')
    plt.ylabel('MSD')
    plt.legend(loc='upper left')
    plt.savefig('MSD_GAS_AB_' + plt_name + '.png', dpi=1000)
    plt.close()

    plt.plot(msd_time, LIQ_A,  color="r", marker='o', markersize=1, linestyle='None', label='Liq_A')
    plt.plot(msd_time, LIQ_B,  color="b", marker='o', markersize=1, linestyle='None', label='Liq_B')
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('Timesteps')
    plt.ylabel('MSD')
    plt.legend(loc='upper left')
    plt.savefig('MSD_LIQ_AB_' + plt_name + '.png', dpi=1000)
    plt.close()

    plt.plot(msd_time, MSD_T,  color="g", marker='o', markersize=1, linestyle='None', label='MSD')
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('Timesteps')
    plt.ylabel('MSD')
    plt.legend(loc='upper left')
    plt.savefig('MSD_total_' + plt_name + '.png', dpi=1000)
    plt.close()

    plt.plot(msd_time, MSD_TL,  color="b", marker='o', markersize=1, linestyle='None', label='Liq')
    plt.plot(msd_time, MSD_TG,  color="r", marker='o', markersize=1, linestyle='None', label='Gas')
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('Timesteps')
    plt.ylabel('MSD')
    plt.legend(loc='upper left')
    plt.savefig('MSD_LG_' + plt_name + '.png', dpi=1000)
    plt.close()

else:                                                           # if monodisperse plot total values
    sns.kdeplot(avg_sys_density[0], shade = True, color="g")
    plt.savefig('avg_density_' + plt_name + '.png', dpi=1000)
    plt.close()

    sns.kdeplot(getDensityPlease(last), shade = True, color="g")
    plt.savefig('final_density_' + plt_name + '.png', dpi=1000)
    plt.close()

    plt.plot(MCS, color="g")
    plt.savefig('MCS_'+ plt_name + '.png', dpi=1000)
    plt.close()

    plt.plot(GF, color="g")
    plt.ylim((0,1))
    plt.savefig('GF_'+plt_name+'.png', dpi=1000)
    plt.close()

    plt.plot(largest, color="g")
    plt.savefig('Largest_clust_'+plt_name+'.png', dpi=1000)
    plt.close()

    plt.plot(msd_time, MSD_T,  color="g", marker='o', markersize=1, linestyle='None', label='MSD')
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('Timesteps')
    plt.ylabel('MSD')
    plt.legend(loc='upper left')
    plt.savefig('MSD_total_' + plt_name + '.png', dpi=1000)
    plt.close()

    plt.plot(msd_time, MSD_TL,  color="b", marker='o', markersize=1, linestyle='None', label='Liq')
    plt.plot(msd_time, MSD_TG,  color="r", marker='o', markersize=1, linestyle='None', label='Gas')
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('Timesteps')
    #plt.xlabel(r'Time ($\tau$)')
    plt.ylabel('MSD')
    plt.legend(loc='upper left')
    plt.savefig('MSD_LG_' + plt_name + '.png', dpi=1000)
    plt.close()


#!/usr/bin/env python
# -*- coding: utf-8 -*-

################################################################################
#
#   RMG - Reaction Mechanism Generator
#
#   Copyright (c) 2002-2009 Prof. William H. Green (whgreen@mit.edu) and the
#   RMG Team (rmg_dev@mit.edu)
#
#   Permission is hereby granted, free of charge, to any person obtaining a
#   copy of this software and associated documentation files (the "Software"),
#   to deal in the Software without restriction, including without limitation
#   the rights to use, copy, modify, merge, publish, distribute, sublicense,
#   and/or sell copies of the Software, and to permit persons to whom the
#   Software is furnished to do so, subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be included in
#   all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
#   THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#   DEALINGS IN THE SOFTWARE.
#
################################################################################

"""
This module contains the :class:`ThermoJob` class, used to compute and save the
thermodynamics information for a single species.
"""

import os.path
import math
import numpy.linalg
import logging

import rmgpy.constants as constants
from rmgpy.cantherm.output import prettify

from rmgpy.statmech.translation import Translation, IdealGasTranslation
from rmgpy.statmech.rotation import Rotation, LinearRotor, NonlinearRotor, KRotor, SphericalTopRotor
from rmgpy.statmech.vibration import Vibration, HarmonicOscillator
from rmgpy.statmech.torsion import Torsion, HinderedRotor
from rmgpy.statmech.conformer import Conformer

from rmgpy.thermo.thermodata import ThermoData
from rmgpy.thermo.nasa import NASAPolynomial, NASA
from rmgpy.thermo.wilhoit import Wilhoit

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

class ThermoJob:
    """
    A representation of a CanTherm thermodynamics job. This job is used to
    compute and save the thermodynamics information for a single species.
    """

    def __init__(self, species, thermoClass):
        self.species = species
        self.thermoClass = thermoClass

    def execute(self, outputFile=None, plot=False):
        """
        Execute the thermodynamics job, saving the results to the
        given `outputFile` on disk.
        """
        self.generateThermo()
        if outputFile is not None:
            self.save(outputFile)
            if plot:
                self.plot(os.path.dirname(outputFile))

    def generateThermo(self):
        """
        Generate the thermodynamic data for the species and fit it to the
        desired heat capacity model (as specified in the `thermoClass`
        attribute).
        """
        if self.thermoClass.lower() not in ['wilhoit', 'nasa']:
            raise Exception('Unknown thermodynamic model "{0}".'.format(self.thermoClass))

        species = self.species

        logging.info('Generating {0} thermo model for {1}...'.format(self.thermoClass, species))

        Tlist = numpy.arange(10.0, 3001.0, 10.0, numpy.float64)
        Cplist = numpy.zeros_like(Tlist)
        H298 = 0.0
        S298 = 0.0
        conformer = self.species.conformer
        for i in range(Tlist.shape[0]):
            Cplist[i] += conformer.getHeatCapacity(Tlist[i])
        H298 += conformer.getEnthalpy(298.) + conformer.E0.value_si
        S298 += conformer.getEntropy(298.)

        if not any([isinstance(mode, (LinearRotor, NonlinearRotor)) for mode in conformer.modes]):
            # Monatomic species
            linear = False
            Nfreq = 0
            Nrotors = 0
            Cp0 = 2.5 * constants.R
            CpInf = 2.5 * constants.R
        else:
            # Polyatomic species
            linear = True if isinstance(conformer.modes[1], LinearRotor) else False
            Nfreq = len(conformer.modes[2].frequencies.value)
            Nrotors = len(conformer.modes[3:])
            Cp0 = (3.5 if linear else 4.0) * constants.R
            CpInf = Cp0 + (Nfreq + 0.5 * Nrotors) * constants.R

        wilhoit = Wilhoit()
        if Nfreq == 0 and Nrotors == 0:
            wilhoit.Cp0 = (Cplist[0],"J/(mol*K)")
            wilhoit.CpInf = (Cplist[0],"J/(mol*K)")
            wilhoit.B = (500.,"K")
            wilhoit.H0 = (0.0,"J/mol")
            wilhoit.S0 = (0.0,"J/(mol*K)")
            wilhoit.H0 =  (H298 -wilhoit.getEnthalpy(298.15), "J/mol")
            wilhoit.S0 = (S298 - wilhoit.getEntropy(298.15),"J/(mol*K)")
        else:
            wilhoit.fitToData(Tlist, Cplist, Cp0, CpInf, H298, S298, B0=500.0)

        if self.thermoClass.lower() == 'nasa':
            species.thermo = wilhoit.toNASA(Tmin=10.0, Tmax=3000.0, Tint=500.0)
        else:
            species.thermo = wilhoit

    def save(self, outputFile):
        """
        Save the results of the thermodynamics job to the file located
        at `path` on disk.
        """
        species = self.species
        logging.info('Saving thermo for {0}...'.format(species.label))

        f = open(outputFile, 'a')

        f.write('# Thermodynamics for {0}:\n'.format(species.label))
        H298 = species.getThermoData().getEnthalpy(298) / 4184.
        S298 = species.getThermoData().getEntropy(298) / 4.184
        f.write('#   Enthalpy of formation (298 K)   = {0:9.3f} kcal/mol\n'.format(H298))
        f.write('#   Entropy of formation (298 K)    = {0:9.3f} cal/(mol*K)\n'.format(S298))
        f.write('#    =========== =========== =========== =========== ===========\n')
        f.write('#    Temperature Heat cap.   Enthalpy    Entropy     Free energy\n')
        f.write('#    (K)         (cal/mol*K) (kcal/mol)  (cal/mol*K) (kcal/mol)\n')
        f.write('#    =========== =========== =========== =========== ===========\n')
        for T in [300,400,500,600,800,1000,1500,2000,2400]:
            Cp = species.getThermoData().getHeatCapacity(T) / 4.184
            H = species.getThermoData().getEnthalpy(T) / 4184.
            S = species.getThermoData().getEntropy(T) / 4.184
            G = species.getThermoData().getFreeEnergy(T) / 4184.
            f.write('#    {0:11g} {1:11.3f} {2:11.3f} {3:11.3f} {4:11.3f}\n'.format(T, Cp, H, S, G))
        f.write('#    =========== =========== =========== =========== ===========\n')

        string = 'thermo(label={0!r}, thermo={1!r})'.format(species.label, species.getThermoData())
        f.write('{0}\n\n'.format(prettify(string)))

        f.close()

        f = open(os.path.join(os.path.dirname(outputFile), 'chem.inp'), 'a')

        thermo = species.getThermoData()
        if isinstance(thermo, NASA):

            poly_low = thermo.polynomials[0]
            poly_high = thermo.polynomials[1]

            # Determine the number of each type of element in the molecule
            elements = ['C','H','N','O']; elementCounts = [0,0,0,0]

            # Remove elements with zero count
            index = 2
            while index < len(elementCounts):
                if elementCounts[index] == 0:
                    del elements[index]
                    del elementCounts[index]
                else:
                    index += 1

            # Line 1
            string = '{0:<16}        '.format(species.label)
            if len(elements) <= 4:
                # Use the original Chemkin syntax for the element counts
                for symbol, count in zip(elements, elementCounts):
                    string += '{0!s:<2}{1:<3d}'.format(symbol, count)
                string += '     ' * (4 - len(elements))
            else:
                string += '     ' * 4
            string += 'G{0:<10.3f}{1:<10.3f}{2:<8.2f}      1'.format(poly_low.Tmin.value_si, poly_high.Tmax.value_si, poly_low.Tmax.value_si)
            if len(elements) > 4:
                string += '&\n'
                # Use the new-style Chemkin syntax for the element counts
                # This will only be recognized by Chemkin 4 or later
                for symbol, count in zip(elements, elementCounts):
                    string += '{0!s:<2}{1:<3d}'.format(symbol, count)
            string += '\n'

            # Line 2
            string += '{0:< 15.8E}{1:< 15.8E}{2:< 15.8E}{3:< 15.8E}{4:< 15.8E}    2\n'.format(poly_high.c0, poly_high.c1, poly_high.c2, poly_high.c3, poly_high.c4)

            # Line 3
            string += '{0:< 15.8E}{1:< 15.8E}{2:< 15.8E}{3:< 15.8E}{4:< 15.8E}    3\n'.format(poly_high.c5, poly_high.c6, poly_low.c0, poly_low.c1, poly_low.c2)

            # Line 4
            string += '{0:< 15.8E}{1:< 15.8E}{2:< 15.8E}{3:< 15.8E}                   4\n'.format(poly_low.c3, poly_low.c4, poly_low.c5, poly_low.c6)

            f.write(string)

            f.close()

    def plot(self, outputDirectory):
        """
        Plot the heat capacity, enthapy, entropy, and Gibbs free energy of the
        fitted thermodynamics model, along with the same values from the
        statistical mechanics model that the thermodynamics model was fitted
        to. The plot is saved to the file ``thermo.pdf`` in the output
        directory. The plot is not generated if ``matplotlib`` is not installed.
        """
        # Skip this step if matplotlib is not installed
        try:
            import pylab
        except ImportError:
            return

        Tlist = numpy.arange(10.0, 2501.0, 10.0)
        Cplist = numpy.zeros_like(Tlist)
        Cplist1 = numpy.zeros_like(Tlist)
        Hlist = numpy.zeros_like(Tlist)
        Hlist1 = numpy.zeros_like(Tlist)
        Slist = numpy.zeros_like(Tlist)
        Slist1 = numpy.zeros_like(Tlist)
        Glist = numpy.zeros_like(Tlist)
        Glist1 = numpy.zeros_like(Tlist)

        conformer = self.species.conformer
        thermo = self.species.getThermoData()
        for i in range(Tlist.shape[0]):
            Cplist[i] = conformer.getHeatCapacity(Tlist[i])
            Slist[i] = conformer.getEntropy(Tlist[i])
            Hlist[i] = (conformer.getEnthalpy(Tlist[i]) + conformer.E0.value_si) * 0.001
            Glist[i] = Hlist[i] - Tlist[i] * Slist[i] * 0.001
            Cplist1[i] = thermo.getHeatCapacity(Tlist[i])
            Slist1[i] = thermo.getEntropy(Tlist[i])
            Hlist1[i] = thermo.getEnthalpy(Tlist[i]) * 0.001
            Glist1[i] = thermo.getFreeEnergy(Tlist[i]) * 0.001

        fig = pylab.figure(figsize=(10,8))

        pylab.subplot(2,2,1)
        pylab.plot(Tlist, Cplist / 4.184, '-r', Tlist, Cplist1 / 4.184, '-b')
        pylab.xlabel('Temperature (K)')
        pylab.ylabel('Heat capacity (cal/mol*K)')
        pylab.legend(['statmech', 'thermo'], loc=4)

        pylab.subplot(2,2,2)
        pylab.plot(Tlist, Slist / 4.184, '-r', Tlist, Slist1 / 4.184, '-b')
        pylab.xlabel('Temperature (K)')
        pylab.ylabel('Entropy (cal/mol*K)')

        pylab.subplot(2,2,3)
        pylab.plot(Tlist, Hlist / 4.184, '-r', Tlist, Hlist1 / 4.184, '-b')
        pylab.xlabel('Temperature (K)')
        pylab.ylabel('Enthalpy (kcal/mol)')

        pylab.subplot(2,2,4)
        pylab.plot(Tlist, Glist / 4.184, '-r', Tlist, Glist1 / 4.184, '-b')
        pylab.xlabel('Temperature (K)')
        pylab.ylabel('Gibbs free energy (kcal/mol)')

        fig.subplots_adjust(left=0.10, bottom=0.08, right=0.95, top=0.95, wspace=0.35, hspace=0.20)
        pylab.savefig(os.path.join(outputDirectory, 'thermo.pdf'))
        pylab.close()


__author__ = 'sebastian'

import os
import copy
import numpy as np
import math
from Utility.Logging_Extension import logger
from Utility.Types.Extrinsics import Extrinsics
from Utility.Types.Intrinsics import Intrinsics
from Utility.Types.Ray import Ray

class Camera(Extrinsics, Intrinsics):

    DEPTH_MAP_WRT_UNIT_VECTORS = "DEPTH_MAP_WRT_UNIT_VECTORS"
    DEPTH_MAP_WRT_CANONICAL_VECTORS = "DEPTH_MAP_WRT_CANONICAL_VECTORS"  # Vectors where the last component is 1

    def __init__(self, file_name=None, width=None, height=None):

        # Coordinate system defines self._quaternion and self.rotation_mat
        super(Camera, self).__init__()

        # Used for visualization
        self.normal = np.array([0, 0, 0], dtype=float)
        # single color to represent the camera (for example in a ply file)
        self.color = np.array([255, 255, 255], dtype=int)

        self.file_name = file_name

        self.width = width
        self.height = height

        # Differentiate between view/image index and reconstruction index
        self.view_index = None          # This is the index w.r.t. to the input images
        self.camera_index = None        # This is the index w.r.t the reconstructed cameras

        self.depth_map_fp = None
        self.depth_map_callback = None
        self.depth_map_semantic = None

    # defines how the class will be printed
    def __repr__(self):
        return self.__str__()

    def __str__(self):
        file_name = self.file_name
        if file_name is None:
            file_name = 'None'
        return str('Camera: ' + file_name + ' ' + str(self._center) + ' ' + str(self.normal))

    def is_monocular_cam(self):
        return True

    # def compare_file_name(self, other):
    #     return self.file_name < other.file_name

    def get_file_name(self):
        return self.file_name

    def get_h5_file_name(self):
        return os.path.splitext(self.file_name)[0] + '.h5'

    def get_viewing_dir_world_coord(self):

        # Positive viewing dir
        viewing_dir_cam_coord = np.asarray([0, 0, 1], dtype=float)
        ray_dir_vec = np.dot(self.get_rotation_mat().T, viewing_dir_cam_coord)
        return ray_dir_vec

    def get_up_dir_world_coord(self):

        # Positive viewing dir
        up_dir_cam_coord = np.asarray([0, 1, 0], dtype=float)
        ray_dir_vec = np.dot(self.get_rotation_mat().T, up_dir_cam_coord)
        return ray_dir_vec

    # ============================= Projecting Single POINTS into Image ==========================================

    def project_single_point_woorld_coord_into_camera_image_as_image_coord(self, point_world_coord):
        """
        $ >>> point_world_coordinates.shape
        $ >>> (3,)
        :param point_world_coord:
        :return:
        """
        image_point, visible = self.project_single_woorld_coord_into_camera_image_as_image_coord(
            point_world_coord.coord)
        return image_point, visible

    # ============================= Projecting Single COORDS into Image ==========================================

    def project_single_woorld_coord_into_camera_image_as_image_coord(self,
                                                                     coord_in_world_coord):
        """
        $ >>> point_world_coordinates.shape
        $ >>> (3,)
        :param coord_in_world_coord:
        :return:
        """
        point_cam_coord = self.world_to_cam_coord_single_coord(
            world_coord=coord_in_world_coord)
        image_point, visible = \
            self.project_single_point_cam_coord_into_camera_image_as_image_coord(
            point_cam_coord)
        return image_point, visible

    def project_single_point_cam_coord_into_camera_image_as_image_coord(self,
                                                                        point_cam_coord):

        """
        THIS RETURNS IMAGE POINTS WITHIN [0, width] x [0, height]
        :param point_cam_coord:
        :param debug_output:
        :return:
        """

        visible = False
        image_point = None

        assert self.width > 0
        assert self.height > 0

        #logger.vinfo('calibration_mat', calibration_mat)

        # To the Visibility of 3D points (on the image plane)
        # Depth of a point with respect to a certain camera is described in Multiple View Geometry on page 162
        # Further Key Words: View Frustum
        # On the other hand: Each point is exactly visible if its projection into the camera coordinate frame
        # has a z component > 0. In addition, we must check that the projection of the point lies within the image
        # bounds (see below)
        if point_cam_coord[2] > 0:

            # Make sure the calibration matrix is valid

            #  project the 3d point onto the 2D camera image plane
            point_image_plane_hom = self.get_calibration_mat().dot(point_cam_coord)

            # See Multiple View Geometry page 155
            # point_image_plane_hom = [fX+Zp_x, fY+Zp_y, Z]
            # with camera point [X,Y,Z] and principal point [p_x, p_y]
            # Remember: two homogeneous coordinates hom_1 and hom_2 with hom_1 = a hom_2 for a scalar 'a' are
            # considered as equivalent
            # Dividing by the last component (i.e. z = 1) results in the image plane coordinates
            # [fX+Zp_x, fY+Zp_y, Z] -> [fX/Z+p_x, fY/Z+p_y, 1]
            point_image_plane_hom = (point_image_plane_hom / point_image_plane_hom[2])

            # x_image_coord and y_image_coord are within [0, width] x [0, height]
            x_image_coord = copy.deepcopy(point_image_plane_hom[0])
            y_image_coord = copy.deepcopy(point_image_plane_hom[1])

            if self.has_radial_distortion():

                # http://ccwu.me/vsfm/doc.html#nvm
                # https://groups.google.com/forum/#!msg/vsfm/IcbdIVv_Uek/Us32SBUNK9oJ

                # NOTE that the parameters saved in NVM file is slightly different with the internal representation.
                # Instead, NVM saves the following for each camera:
                # f, R (as quaternion), C = - R'T, rn = r * f * f.
                # The PBA code includes functions for loading the NVM file and convert the camera parameters.
                rn = self.get_radial_distortion()

                r = rn / float(self.get_focal_length() * self.get_focal_length())

                # mx and my are centered around the principal point
                principal_point = self.get_principal_point()
                mx = x_image_coord - principal_point[0]
                my = y_image_coord - principal_point[1]

                # The distortion factor is r2 = r * (mx * mx + my * my)
                r_2 = r * (mx * mx + my * my)
                logger.vinfo('r_2', r_2)

                # The undistorted measurement is (1 + r2) * (mx, my)
                undistorted_x = (1 + r_2) * mx
                undistorted_y = (1 + r_2) * my

                logger.vinfo('x_image_coord (without undistortion)', x_image_coord)
                logger.vinfo('y_image_coord (without undistortion)', y_image_coord)

                x_image_coord = undistorted_x + principal_point[0]
                y_image_coord = undistorted_y + principal_point[1]
                logger.vinfo('x_image_coord (with undistortion)', x_image_coord)
                logger.vinfo('y_image_coord (with undistortion)', y_image_coord)

                # THE UNDISTORTED VALUES ARE NOT REALLY BETTER THAN THE ORIGINAL ONES
                # TODO MAYBE BUG IS FIXED NOW
                # TODO INVESTIGATE
                # TODO TEST WITH FISHEYE
                assert False    # Radial Distortion support is disabled for now

            on_image_plane_x = abs(x_image_coord) < self.width
            on_image_plane_y = abs(y_image_coord) < self.height

            visible = on_image_plane_x and on_image_plane_y

            image_point = (x_image_coord, y_image_coord)

        return image_point, visible

    def compute_reprojection_error_single_point(self, point):
        projected_image_point, visible = \
            self.project_single_woorld_coord_into_camera_image_as_image_coord(
                point.coord)
        if not visible:
            return None
        for measurement in point.measurements:
            if measurement.camera_index == self.camera_index:
                error_x = measurement.x - projected_image_point[0]
                error_y = measurement.y - projected_image_point[1]
                error = math.sqrt(error_x * error_x + error_y * error_y)
                return error

    # ============================= Projecting Multiple POINTS into Image ==========================================

    def project_multiple_points_woorld_coord_into_camera_image_as_image_coord(self, points_world_coord):
        """
        $ >>> point_world_coordinates.shape
        $ >>> (3,)
        :param point_world_coord:
        :return:
        """

        image_points = []
        visibility_flags = []
        for point in points_world_coord:
            image_point, visible = \
                self.project_single_point_woorld_coord_into_camera_image_as_image_coord(point)
            image_points.append(image_point)
            visibility_flags.append(visible)
        return image_points, visibility_flags

    # ============================= Projecting Multiple COORDS into Image ==========================================

    def project_multiple_world_coords_into_camera_image_as_image_coord(self,
                                                                       coords_world_coord):
        """
        :param coords_world_coord: is a single array
        :return: points_image_plane: is not necessarily lying on the image plane (check visibility array)
        :return: visibility_array: contains True and False values (visibility_array.dtype is bool)
        """

        #logger.info('project_multiple_world_coords_into_camera_image_as_image_coord: ...')

        # shape = rows, columns
        #logger.info('points_world_coordinates.shape: ' + str(coords_world_coord.shape))
        assert coords_world_coord.shape[0] == 3 and coords_world_coord.shape[1] > 3
        assert self.width is not None
        assert self.height is not None

        if self.has_radial_distortion():
            assert False    # Radial distortion not supported atm

        # transformation of the object coordinate system to the camera coordinate system
        # point_cam_coord = (self.get_rotation_mat()).dot(point_world_coordinates - self.get_camera_center())

        # make coordinate homogeneous and apply transformation
        hom_row = np.ones(coords_world_coord.shape[1])
        # can't use hstack or vstack here, since dimensions do not match
        points_world_coordinates_hom = np.row_stack((coords_world_coord, hom_row))
        points_cam_coord_hom = self.get_4x4_world_to_cam_mat().dot(
            points_world_coordinates_hom)

        coord_cam_coord = points_cam_coord_hom[0:3]
        # Rows first, columns later => points_world_coordinates[row_index][column_index]
        # points_world_coordinates has 3 ROWS and N COLUMNS (N >> 3)
        # The first/second/third row contains the x/y/z values for ALL points
        # points_world_coordinates[row_index] is a np.array with N elements
        #logger.vinfo('points_world_coordinates.shape', points_world_coordinates.shape)
        #logger.vinfo('coord_cam_coord.shape', coord_cam_coord.shape)

        # logger.debug('============DEBUG============')
        # points_cam_coord_ref = np.empty([3, 0], dtype=float)
        # for ground_point in ground_points:
        #     point_world_coordinates = ground_point.get_coord_as_array()
        #     point_cam_coordinates = self.world_to_cam_coord_single_point(point_world_coordinates)
        #     points_cam_coord_ref = np.column_stack((points_cam_coord_ref, point_cam_coordinates))
        # # 'allclose' test due to numerical differences
        # comp_val = np.allclose(coord_cam_coord, points_cam_coord_ref)
        # assert comp_val
        # logger.debug('=============================')

        points_image_plane, visibility_array = \
            self.project_multiple_cam_coords_into_camera_image_as_image_coord(
                coord_cam_coord)

        #logger.info('project_multiple_world_coords_into_camera_image_as_image_coord: Done')

        return points_image_plane, visibility_array

    def project_multiple_cam_coords_into_camera_image_as_image_coord(self, coords_cam_coord):
        """

        :param coords_cam_coord:
        :return: points_image_plane: is not necessarily lying on the image plane (check visibility_array)
        :return: visibility_array: contains True and False values (visibility_array.dtype is bool)
        """

        #logger.vinfo('points_cam_coord.shape', coords_cam_coord.shape)
        assert coords_cam_coord.shape[0] == 3 and coords_cam_coord.shape[1] > 3

        points_image_plane_hom = self.get_calibration_mat().dot(coords_cam_coord)
        #points_image_plane_hom = (self.get_calibration_mat()).dot(coords_cam_coord)
        # logger.vinfo('points_image_plane_hom.shape', points_image_plane_hom.shape)

        # element wise ops: np.true_divide(), np.greater()
        in_front_of_array = np.greater(coords_cam_coord[2], np.zeros(len(coords_cam_coord[2]), dtype=float))

        # normalization of the homogeneous coordinates
        points_image_plane_hom[0] = np.true_divide(points_image_plane_hom[0], points_image_plane_hom[2])
        points_image_plane_hom[1] = np.true_divide(points_image_plane_hom[1], points_image_plane_hom[2])

        # use the absolute value to safe comparisons
        on_image_plane_x = np.logical_and(points_image_plane_hom[0] >= 0, points_image_plane_hom[0] < self.width)
        on_image_plane_y = np.logical_and(points_image_plane_hom[1] >= 0, points_image_plane_hom[1] < self.height)

        # np.logical_and does not support 3 arguments (3 parameter is the output array)
        visibility_array = np.logical_and(in_front_of_array, on_image_plane_x)
        visibility_array = np.logical_and(visibility_array, on_image_plane_y)

        points_image_plane_hom = copy.deepcopy(points_image_plane_hom)

        return points_image_plane_hom[0:2].T, visibility_array

    # ============================= Check Visibility ==========================================

    def check_visibilty_of_single_point_world_coords(self, point_world_coordinates):
        _, visible = self.project_single_point_woorld_coord_into_camera_image_as_image_coord(
            point_world_coordinates)
        return visible

    def check_visibilty_of_single_point_cam_coords(self, point_cam_coord):
        _, visible = self.project_single_point_cam_coord_into_camera_image_as_image_coord(
            point_cam_coord)
        return visible

    def check_visibility_of_points_world_coords(self, points_world_coordinates):
        """
        Checks the visibility of points in camera coordinates. Returns a binary array
        $ >>> points_world_coordinates.shape
        $ >>> (n,3)                 # where n is the number of points
        :param points_world_coordinates:
        :return:
        """
        logger.info('check_visibility_of_world_points: ...')
        _, visibility_array = \
            self.project_multiple_world_coords_into_camera_image_as_image_coord(
                points_world_coordinates)
        logger.info('check_visibility_of_world_points: Done')
        return visibility_array

    def check_visibility_of_points_cam_coords(self, points_cam_coord):
        """
        Checks the visibility of points in camera coordinates. Returns a binary array
        $ >>> points_world_coordinates.shape
        $ >>> (n,3)                 # where n is the number of points
        """
        logger.info('check_visibility_of_cam_points: ...' )
        _, visibility_array = \
            self.project_multiple_cam_coords_into_camera_image_as_image_coord(points_cam_coord)
        logger.info('check_visibility_of_cam_points: Done')
        return visibility_array

    # ============================== Measurements ===============================================================
    def compute_measurement_pos_of_points(self, points):

        measurement_pos = []
        for point in points:
            for measurement in point.measurements:
                if measurement.camera_index == self.camera_index:
                    measurement_pos.append(measurement.get_x_y())
        return measurement_pos

    # ============================== Camera Rays ===============================================================

    def generate_camera_rays(self, x_y_positions, convert_to_world_coords=True):
        """
        SUPPORTS ATM ONLY "SIMPLE_PINHOLE", i.e. f,cx,cy and a scale ratio of 0
        """

        assert self.width is not None
        assert self.height is not None

        focal_length = self._calibration_mat[0][0]
        rays = []
        for x_y in x_y_positions:
            #logger.info(x_y)
            assert x_y[0] >= 0 and x_y[1] >= 0

            # Compute the direction vector in CAMERA COORDINATES
            ray_dir_vec = np.array([x_y[0] - self.width/float(2),       # x
                                    x_y[1] - self.height/float(2),      # y
                                    focal_length],
                                   dtype=float)
            ray_pos_vec = np.zeros(3)

            if convert_to_world_coords:
                # Compute the direction vector in WORLD COORDINATES
                ray_dir_vec = np.dot(self.get_rotation_mat().T, ray_dir_vec)
                ray_pos_vec = self.get_camera_center()

            rays.append(Ray(pos_vec=ray_pos_vec, dir_vec=ray_dir_vec))
        return rays

    def set_depth_map(self, depth_map_ifp, depth_map_callback, depth_map_semantic):
        self.depth_map_fp = depth_map_ifp
        self.depth_map_callback = depth_map_callback
        self.depth_map_semantic = depth_map_semantic

    def get_depth_map(self):
        if os.path.isfile(self.depth_map_fp):
            return self.depth_map_callback(self.depth_map_fp)
        else:
            return None

    def convert_depth_map_to_world_coords(self,
                                          depth_map,
                                          depth_map_semantic,
                                          shift_to_pixel_center,        # False for Colmap, True for MVE
                                          depth_map_display_sparsity=100,
                                          inverted_cam_model=False):
        """
        Do not confuse z_buffer with depth_buffer!
        z_buffer contains values in [0,1]
        depth_buffer contains the actual distance values

        :param depth_buffer_matrix:
        :param n_th_result_point:
        :return:
        """
        logger.info('Converting depth map to world coordinates: ...')
        cam_coords = self.convert_depth_map_to_cam_coords(
            depth_map,
            depth_map_semantic,
            shift_to_pixel_center,
            depth_map_display_sparsity,
            inverted_cam_model=inverted_cam_model)

        world_coords = self.cam_to_world_coord_multiple_coords(
            cam_coords)

        logger.info('Converting depth map to world coordinates: Done')
        return world_coords

    def convert_depth_map_to_cam_coords(self,
                                        depth_map,
                                        depth_map_semantic,
                                        shift_to_pixel_center,  # False for Colmap, True for MVE
                                        depth_map_display_sparsity=100,
                                        inverted_cam_model=False):

        assert 0 < depth_map_display_sparsity

        height, width = depth_map.shape
        logger.info('height ' + str(height))
        logger.info('width ' + str(width))

        if self.height == height and self.width == width:
            x_step_size = 1.0
            y_step_size = 1.0
        else:
            x_step_size = self.width / width
            y_step_size = self.height / height
            logger.info('x_step_size ' + str(x_step_size))
            logger.info('y_step_size ' + str(y_step_size))

        fx, fy, skew, cx, cy = self.split_intrinsic_mat(self.get_calibration_mat())
        logger.vinfo('fx, fy, skew, cx, cy: ', str([fx, fy, skew, cx, cy]))

        indices = np.indices((height, width))
        y_index_list = indices[0].flatten()
        x_index_list = indices[1].flatten()

        if inverted_cam_model:  # For Blender, VTK, etc
            # Use the local coordinate system of the camera to analyze its viewing directions
            # The Blender camera coordinate system looks along the negative z axis (blue),
            # the up axis points along the y axis (green).
            y_index_list = y_index_list[::-1]  # Reverse order of indices
            x_index_list = x_index_list[::-1]  # Reverse order of indices
            fx = -fx
            fy = -fy
            assert False    # TODO Verify this

        depth_values = depth_map.flatten()

        assert len(x_index_list) == len(y_index_list) == len(depth_values)

        if shift_to_pixel_center:
            # https://github.com/simonfuhrmann/mve/blob/master/libs/mve/depthmap.cc
            #  math::Vec3f v = invproj * math::Vec3f(
            #       (float)x + 0.5f, (float)y + 0.5f, 1.0f);
            u_index_coord_list = x_step_size * x_index_list + 0.5
            v_index_coord_list = y_step_size * y_index_list + 0.5
        else:
            # https://github.com/colmap/colmap/blob/dev/src/base/reconstruction.cc
            #   // COLMAP assumes that the upper left pixel center is (0.5, 0.5)
            # i.e. pixels are already shifted
            u_index_coord_list = x_step_size * x_index_list
            v_index_coord_list = y_step_size * y_index_list

        # The cannoncial vectors are defined according to p.155 of
        # "Multiple View Geometry" by Hartley and Zisserman using a canonical
        # focal length of 1 , i.e. vec = [(x - cx) / fx, (y - cy) / fy, 1]
        x_coords_canonical = (u_index_coord_list - cx) / fx + (cy - v_index_coord_list) * skew / (fx * fy)
        y_coords_canonical = (v_index_coord_list - cy) / fy
        z_coords_canonical = np.ones(len(depth_values), dtype=float)

        # Determine non-background data
        # non_background_flags = np.logical_not(np.isnan(depth_values))
        depth_values_not_nan = np.nan_to_num(depth_values)
        non_background_flags = depth_values_not_nan > 0

        x_coords_canonical_filtered = x_coords_canonical[non_background_flags]
        y_coords_canonical_filtered = y_coords_canonical[non_background_flags]
        z_coords_canonical_filtered = z_coords_canonical[non_background_flags]
        depth_values_filtered = depth_values[non_background_flags]

        if depth_map_display_sparsity != 100:
            x_coords_canonical_filtered = x_coords_canonical_filtered[::depth_map_display_sparsity]
            y_coords_canonical_filtered = y_coords_canonical_filtered[::depth_map_display_sparsity]
            z_coords_canonical_filtered = z_coords_canonical_filtered[::depth_map_display_sparsity]
            depth_values_filtered = depth_values_filtered[::depth_map_display_sparsity]

        if depth_map_semantic == Camera.DEPTH_MAP_WRT_CANONICAL_VECTORS:
            # In this case, the depth values are defined w.r.t. the canonical
            # vectors. This kind of depth data is used by Colmap.
            x_coords_filtered = x_coords_canonical_filtered * depth_values_filtered
            y_coords_filtered = y_coords_canonical_filtered * depth_values_filtered
            z_coords_filtered = z_coords_canonical_filtered * depth_values_filtered

        elif depth_map_semantic == Camera.DEPTH_MAP_WRT_UNIT_VECTORS:
            # In this case the depth values are defined w.r.t. the normalized
            # canonical vectors. This kind of depth data is used by MVE.
            cannonical_norms_filtered = np.linalg.norm(
                np.array(
                    [x_coords_canonical_filtered,
                     y_coords_canonical_filtered,
                     z_coords_canonical_filtered],
                    dtype=float),
                axis=0)
            # Instead of normalizing the x,y and z component, we divide the
            # depth values by the corresponding norm.
            normalized_depth_values_filtered = depth_values_filtered / cannonical_norms_filtered
            x_coords_filtered = x_coords_canonical_filtered * normalized_depth_values_filtered
            y_coords_filtered = y_coords_canonical_filtered * normalized_depth_values_filtered
            z_coords_filtered = z_coords_canonical_filtered * normalized_depth_values_filtered

        else:
            assert False

        cam_coords = np.dstack(
            (x_coords_filtered,
             y_coords_filtered,
             z_coords_filtered))[0]

        return cam_coords

    @staticmethod
    def parse_camera_image_files(cameras, path_to_images):
        from PIL import Image
        for camera in cameras:
            # this does NOT load the data; into memory -> should be fast!
            image = Image.open(os.path.join(path_to_images, camera.file_name))
            camera.width, camera.height = image.size
        return cameras

    @staticmethod
    def parse_camera_h5_files(cameras, path_to_h5_files):
        from Utility.File_Handler.H5_File_Handler import H5FileHandler
        for camera in cameras:
            seg = H5FileHandler.read_h5(os.path.join(path_to_h5_files, camera.get_h5_file_name()))
            camera.height, camera.width = seg.shape
        return cameras

    # ============================== Depth Buffer Legacy =============================================================

    # def convert_depth_buffer_to_cam_coords_legazy(self,
    #                                               depth_buffer_matrix,
    #                                               depth_scale_value=1.0,
    #                                               background_depth_value=1000,
    #                                               invert_y=False,
    #                                               n_th_result_point=100
    #                                               ):
    #     """
    #     Do not confuse z_buffer with depth_buffer!
    #     z_buffer contains values in [0,1]
    #     depth_buffer contains the actual distance values
    #
    #     Computes only 1/n_th_result_point points for the point cloud
    #     :param depth_buffer_matrix:
    #     :param n_th_result_point:
    #     :return:
    #     """
    #
    #     logger.info('convert_depth_matrix_to_camera_coords: ...')
    #
    #     assert self.height is not None
    #     assert self.width is not None
    #
    #     depth_buffer_matrix *= depth_scale_value
    #
    #     fx = self.get_calibration_mat()[0][0]
    #     logger.vinfo('focal_length_in_pixel', fx)
    #
    #     # use the local coordinate system of the camera
    #     # to analyze its viewing directions
    #     # blender camera coordinate system looks along the negative z axis (blue)
    #     # the up axis points along the y axis (green)
    #
    #     coords_cam_coord = []
    #     result_point_idx = 0
    #     num_pixels = self.width * self.height
    #     for (y, x), depth_value in np.ndenumerate(depth_buffer_matrix):
    #
    #         if depth_value < background_depth_value:
    #             if result_point_idx % n_th_result_point == 0:
    #                 if result_point_idx % 10000 == 0:
    #                     logger.info('result_point_idx ' + str(result_point_idx) + ' ' + str(num_pixels))
    #
    #                 cx, cy = self.get_principal_point()
    #
    #                 x = self.width - x  # Invert x
    #                 if invert_y:
    #                     y = self.height - y
    #
    #                 x = x - cx  # Shift x from image to camera coordinates
    #                 y = y - cy  # Shift y from image to camera coordinates
    #                 z = -fx  # Invert z
    #                 pixel_vec = np.array([x, y, z], dtype=float)
    #
    #                 # convert to inhomogeneous coordinates
    #                 # divide using the focal length (in pixels)
    #                 # (i.e. the resulting vector is unit free)
    #                 pixel_vec /= pixel_vec[2]
    #                 # multiply by depth value in blender units / meters
    #                 # (i.e. the resulting vector is in blender units / meters)
    #                 pixel_vec *= depth_value
    #                 coords_cam_coord.append(pixel_vec)
    #             result_point_idx += 1
    #
    #     logger.info('convert_depth_matrix_to_camera_coords: Done')
    #     return coords_cam_coord

    # def convert_depth_buffer_to_world_coords_legazy(self,
    #                                                 depth_buffer_matrix,
    #                                                 depth_scale_value=1.0,
    #                                                 background_depth_value=1000,
    #                                                 invert_y=False,
    #                                                 n_th_result_point=100
    #                                                 ):
    #
    #     coords_cam_coord = self.convert_depth_buffer_to_cam_coords_legazy(
    #         depth_buffer_matrix=depth_buffer_matrix,
    #         depth_scale_value=depth_scale_value,
    #         background_depth_value=background_depth_value,
    #         invert_y=invert_y,
    #         n_th_result_point=n_th_result_point
    #     )
    #
    #     coords_world_coord = self.cam_to_world_coord_multiple_coords(
    #         coords_cam_coord)
    #     return coords_world_coord