# -*- coding: utf-8 -*-

import struct
import numpy as np
import multiprocessing # To get CPU count.
import os
import tempfile
import subprocess

import phstuff.barcode as bc

"""Handle marshalling data into and out of DIPHA's formats, and
optionally manage running the DIPHA executable.

"""

# FIXME, TODO: These magic numbers should really be read from the
# DIPHA header files.
DIPHA_MAGIC = 8067171840
DIPHA_WEIGHTED_BOUNDARY_MATRIX = 0
DIPHA_IMAGE_DATA = 1
DIPHA_PERSISTENCE_DIAGRAM = 2
DIPHA_DISTANCE_MATRIX = 7
DIPHA_SPARSE_DISTANCE_MATRIX = 8


def save_weight_matrix(fname, weights):
    """Write a NumPy weight matrix to a DIPHA full distance matrix
    file. Attention: DIPHA uses a convention wherein an edge's
    filtration value is *half* its weight. This function compensates
    by multiplying every weight by 2. Pay attention in case DIPHA's
    behavior changes in the future!

    Parameters:
    -----------

    fname: Name of file to write.

    weights: NumPy array. Elements will be converted to IEEE-754
       doubles and used as weights. The entire array is passed on to
       DIPHA, whose behvior is undefined if it's not symmetric.

    """

    # override_dipha_half: DIPHA currently (commit 0081862) divides all
    #    entries by 2 when loading the file. If this argument is `True`,
    #    we compensate by multiplying all entries by 2. Be sure to pay
    #    attention if DIPHA's behavior changes!

    factor = 2.0
    
    m = weights.shape[0]
    if weights.shape[0] != weights.shape[1]:
        raise ValueError("Matrix is not square.")

    with open(fname, "wb") as f:
        f.write(struct.pack("<q", DIPHA_MAGIC))
        f.write(struct.pack("<q", DIPHA_DISTANCE_MATRIX))
        f.write(struct.pack("<q", m))

        # Why write in a loop? I heard rumors that struct.pack will
        # allocate a lot of temporary heap space if given lots of
        # data. I haven't actually tested.
        for i in range(0, m):
            for j in range(0, m):
                f.write(struct.pack("<d", factor*weights[i,j]))

def save_masked_weight_matrix(fname, weights):
    """Write a masked NumPy weight matrix to a DIPHA sparse distance
    matrix file, keeping edges only for unmasked entries. Attention:
    DIPHA uses a convention wherein an edge's filtration value is
    *half* its weight. This function compensates by multiplying every
    weight by 2. Pay attention in case DIPHA's behavior changes in the
    future!

    Parameters:
    -----------

    fname: Name of file to write.

    weights: Masked NumPy array. Unmasked elements will be converted
       to IEEE-754 and used as weights. Must be symmetric (not
       checked).

    """

    factor = 2.0

    n = weights.shape[1]

    with open(fname, "wb") as f:
        f.write(struct.pack("<q", DIPHA_MAGIC))
        f.write(struct.pack("<q", DIPHA_SPARSE_DISTANCE_MATRIX))
        f.write(struct.pack("<q", n))

        for i in range(0, n):
            f.write(struct.pack('<q', n - np.sum(weights.mask[i, :])))
        for i in range(0, n):
            for j in np.arange(0, n)[-(weights.mask[i, :])]:
                    f.write(struct.pack('<q', j))
        for i in range(0, n):
            for j in np.arange(0, n)[-(weights.mask[i, :])]:
                    f.write(struct.pack('<d', factor*weights[i, j]))
    


def save_edge_list(fname, edge_list):
    """Write a possibly non-complete weighted graph represented as an edge
    list to a DIPHA sparse distance file. Attention: DIPHA uses a
    convention wherein an edge's filtration value is *half* its
    weight. This function compensates by multiplying every weight by
    2. Pay attention in case DIPHA's behavior changes in the future!

    Parameters:
    -----------

    fname: Name of file to write.

    edge_list: A list of the edges and weights of each node, in node
       order. If node `i` has edges to nodes `a_{i,0}, …¸ a_{i,k_i}`
       with weights `w_{i,0}, …, w_{i,k_i}`, then this argument should
       be

          `[[(a_{0,0}, w_{0,0}), …, (a_{0,k_0}, w_{0,k_0})],
            [(a_{1,0}, w_{1,0}), …, (a_{1,k_1}, w_{1,k_1})], 
            …, 
            [(a_{n,0}, w_{n,0}), …, (a_{n,k_n}, w_{n,k_n})]]`

       Remember that an isolated node corresponds to an empty (inner)
       list.

    """

    factor = 2.0
    
    n = len(edge_list)

    with open(fname, "wb") as f:
        f.write(struct.pack("<q", DIPHA_MAGIC))
        f.write(struct.pack("<q", DIPHA_SPARSE_DISTANCE_MATRIX))
        f.write(struct.pack("<q", n))

        for i in range(0, n):
            f.write(struct.pack('<q', len(edge_list[i])))
        for i in range(0, n):
            for j in range(0, len(edge_list[i])):
                f.write(struct.pack('<q', edge_list[i][j][0]))
        for i in range(0, n):
            for j in range(0, len(edge_list[i])):
                f.write(struct.pack('<d', factor*edge_list[i][j][1]))

def save_cubical(fname, complex):
    """Save a (any-dimensional) array that will be interpreted as a
    cubical complex with top-dimensional cells entering the filtration
    at the scale given by the array values.

    Parameters:
    -----------
    
    fname: Name of file to write.

    complex: Any-dimensional array with floating-point values giving
    the filtration values of the top-dimensional cells.
    
    """

    with open(fname, "wb") as f:
        f.write(struct.pack("<q", DIPHA_MAGIC))
        f.write(struct.pack("<q", DIPHA_IMAGE_DATA))
        f.write(struct.pack("<q", np.product(complex.shape)))
        f.write(struct.pack("<q", len(complex.shape)))
        
        for n in complex.shape:
            f.write(struct.pack("<q", n))

        for w in complex.flat:
            f.write(struct.pack("<d", w))
            

def load_barcode(fname, top_dim = None):

    """Load a barcode as written by DIPHA.

    Parameters:
    -----------

    fname: File name to load.
    
    top_dim: Integer. Exclude degrees above this.

    Returns:
    --------

    A dict keyed on persistent homology degree. Associated to each
    degree is a list of `Interval`s.

    """
    
    ret = dict()
    
    with open(fname, "rb") as f:
        if struct.unpack('<q', f.read(8))[0] != DIPHA_MAGIC:
            raise IOError("File %s is not a valid DIPHA file." %(fname))
        if struct.unpack('<q', f.read(8))[0] != DIPHA_PERSISTENCE_DIAGRAM:
            raise IOError("File %s is not a valid DIPHA barcode file." %(fname))

        n = struct.unpack('<q', f.read(8))[0]
        
        for i in range(0, n):
            (d, birth, death) = struct.unpack('<qdd', f.read(3*8))
            if d < 0:
                dim = -d -1
            else:
                dim = d

            if top_dim is None or dim <= top_dim:
                if d < 0:
                    ret.setdefault(dim, []).append(bc.Interval(birth))
                else:
                    ret.setdefault(dim, []).append(bc.Interval(birth, death))

    return ret

def save_text_barcode(fname, barcode):
    """Saves a barcode as a human- and machine-readable text file in an
    ad-hoc format.

    Parameters:
    -----------
    
    fname: File name to save to.
    
    barcode: The barcode to save, a dictionary keyed on degree, each
    entry being a list of `Interval`s.

    """

    with open(fname, "w") as f:
        for dim in barcode.keys():
            for interval in barcode[dim]:
                if interval.is_finite():
                    f.write("%d %g %g\n" %(dim, interval.birth, interval.death))
                else:
                    f.write("%d %g inf\n" %(dim, interval.birth))

    

class DiphaRunner:
    """For long-running computations that may fail, or computations
    involving huge amounts of data, I recommend that you run DIPHA
    yourself. If you prefer to avoid doing that, this class helps you
    run DIPHA without leaving your Python script.

    In order to do so, the `DiphaRunner` object needs to know where to
    find the `mpirun` and `dipha` executables. If they are somewhere
    in your environment's `PATH`, then everything is fine. If not, you
    need to specify where these can be foind either by setting the
    environment variables `DIPHA`/`MPIRUN` (to the full path of the
    executables, not just their directories) or through the optional
    arguments of the class constructor.

    """

    def __init__(self, top_dim, cpu_count = None, quiet = True, mpirun = None, dipha = None):
        """Constructor.
        
        Parameters:
        -----------
        
        top_dim: Highest simplex dimension to include.

        cpu_count (optional): The number of parallell DIPHA processes
        to run. If `None` (default), try to autodetect. The actual
        number that will be used can be found in the `cpu_count`
        member of the returned object.

        quiet (optional): Don't print anything in normal situations
        (default).

        mpirun (optional): Full path of the `mpirun` executable. If
        `None` (default), use environment variable `MPIRUN` or look in
        `PATH`.

        dipha (optional): Full path of the `dipha` executable. If
        `None` (default), use environment variable `DIPHA` or look in
        `PATH`.

        """
    
        self.top_dim = int(top_dim)
        if top_dim <= 0:
            raise ValueError("Top dimension must be postitive.")
        
        if cpu_count is None:
            try:
                self.cpu_count = multiprocessing.cpu_count()
            except NotImplementedError as e:
                self.cpu_count = 1
        else:
            self.cpu_count = cpu_count

        self.quiet = quiet

        if mpirun is None:
            self.mpirun = os.environ.get("MPIRUN")
            paths = os.environ.get("PATH", "")
            if self.mpirun is None:
                for path in paths.split(":"):
                    x = os.path.join(path, "mpirun")
                    if os.path.isfile(x):
                        self.mpirun = x
                        break
        else:
            self.mpirun = mpirun

        if dipha is None:
            self.dipha = os.environ.get("DIPHA")
            paths = os.environ.get("PATH", "")
            if self.dipha is None:
                for path in paths.split(":"):
                    x = os.path.join(path, "dipha")
                    if os.path.isfile(x):
                        self.dipha = x
                        break
        else:
            self.dipha = dipha

        if self.mpirun is None or not os.path.isfile(self.mpirun):
            raise RuntimeError("Could not find mpirun exectuable.")
        if self.dipha is None or not os.path.isfile(self.dipha):
            raise RuntimeError("Could not find dipha executable.")

        self.tmpdir = tempfile.mkdtemp()
        self.ifile = os.path.join(self.tmpdir, "input.dipha")
        self.ofile = os.path.join(self.tmpdir, "output.dipha")
        self.ready = False

        
    def weight_matrix(self, weights):
        """Initialize to run with a weight matrix. See `save_weight_matrix`."""
        
        save_weight_matrix(self.ifile, weights)
        self.ready = True

    def masked_weight_matrix(self, weights):
        """Initialize to run with a masked weight matrix. See
        `save_masked_weight_matrix`."""
        
        save_masked_weight_matrix(self.ifile, weights)
        self.ready = True

    def edge_list(self, edge_list):
        """Initialize to run with an edge list. See `save_edge_list`."""
        
        save_edge_list(self.ifile, edge_list)
        self.ready = True

    def run(self):
        """Run DIPHA. Success is indicated by `code` member of the object
        being zero. If that is the case, then the result is available
        in the `barcode` member."""

        if not self.ready:
            self.cleanup()
            raise RuntimeError("Not ready to run. Maybe you haven't given me data, or I've already been run?")
        
        if self.quiet:
            opipe = subprocess.PIPE
        else:
            opipe = None

        if not self.quiet:
            print("Spawning %d DIPHA processes." %(self.cpu_count))
            
        proc = subprocess.Popen([self.mpirun, "-np", "%d" %(self.cpu_count), self.dipha, "--upper_dim", "%d" %(self.top_dim), self.ifile, self.ofile], stdout=opipe)

        self.out_message = proc.communicate()[0]
        self.code = proc.returncode
        
        if not self.quiet and self.code != 0:
            print("DIPHA exited abnormally with code %d." %(self.code))
            self.cleanup()

        if self.code == 0:
            self.barcode = load_barcode(self.ofile, top_dim = self.top_dim - 1)
            if not self.quiet:
                print("DIPHA exited normally. Result available.")

        self.ready = False
        self.cleanup()
        

    def cleanup(self):
        os.remove(self.ifile)
        os.remove(self.ofile)
        os.rmdir(self.tmpdir)
        self.ifile = None
        self.ofile = None
        self.tmpdir = None