
import numpy as np
import time
import csv
import multiprocessing as mp
import shutil

from cryoem.data.griddata import GridDataList, GridDataReal
from cryoem.model.prior import TransformationPrior
from cryoem.model.prior import TransformationPriors
from cryoem.model.prior import IsoMixturePrior3D, Gamma, Dirichlet
from cryoem.model.prior import IsoTransformPrior
from cryoem.model.gaussian import IsoGaussianND
from cryoem.sampling.isotransformsampler import IsoTransformSampler
from cryoem.sampling.global_rotation_sampler import (
    BruteForceRotationSampler)
from cryoem.data.grid import Grids
from cryoem.model.isomixture import IsoMixture3D
from cryoem.model.transformation import Transformations
from cryoem.program.utils import read_csv_dict, format_time

        


class Rec(object):
    """Run reconstruction."""

    def __init__(self, setup):
        """
        Args:
            setup: RecSetup object, parameters for reconstruction.
        """

        self.setup = setup

    def construct_prior(self, num_projections, grid2D):
        """Return prior over mixture model and transformations."""

        K = self.setup.num_components

        if self.setup.mean_variance:
            mean_variance = self.setup.mean_variance
        else:
            # Estimate mean variance for input data.
            # Ensure that projections of initial models fit inside grid.
            s = grid2D.shape[0] * grid2D.spacing[0] # side length
            mean_variance = (s / 6)**2
        if self.setup.verbose >= 2:
            print('Prior mean std = %1.3f.' % np.sqrt(mean_variance))

        mean_prior = IsoGaussianND(np.zeros(3), mean_variance)
        precision_prior = Gamma(self.setup.gamma_alpha,
                                self.setup.gamma_beta)
        weights_prior = Dirichlet(np.ones(K) * self.setup.alpha_0)

        isomixture_prior = IsoMixturePrior3D(mean_prior,
                                             precision_prior,
                                             weights_prior)

        transformations_prior = [TransformationPriors.create_uniform()
                                 for i in range(num_projections)]

        prior = IsoTransformPrior(isomixture_prior,
                                  transformations_prior)

        return prior

    def create_initial_model(self, prior):
        """Create starting model for Gibbs sampler.

        Model is sampled from prior unless user specified initial
        mixture or transformations.
        """

        model = prior.sample()

        if self.setup.initial_mixture_file:
            mixture = IsoMixture3D.load_csv(
                self.setup.initial_mixture_file)
            model.set_mixture(mixture)
            if self.setup.verbose >= 2:
                print('Initial mixture loaded from file.')
        else:
            if self.setup.verbose >= 2:
                print('Initial mixture sampled from prior.')

        if self.setup.initial_transformations_file:
            ts = Transformations.load3D_csv(
                self.setup.initial_transformations_file)
            model.set_transformations(ts.to_list())
            if self.setup.verbose >= 2:
                print('Initial transformations loaded from file.')
        else:
            if self.setup.verbose >= 2:
                print('Initial transformations set to zero.')



        return model

    def prepare_samplers(self, sampler, global_sampler):
        """Return list of functions of model.

        Each function is a single sampling step.
        """

        # Get required parameters.
        num_steps = self.setup.num_steps
        num_mixture_steps = self.setup.num_mixture_steps
        num_transformation_steps = self.setup.num_transformation_steps

        if not self.setup.use_global:
            if self.setup.verbose >= 2:
                print('Global rotation sampling deactivated.')
            method = sampler.sample
            fix_m = self.setup.fix_mixture
            fix_t = self.setup.fix_transformations
            if fix_t and not fix_m:
                method = sampler.sample_mixture_part
                if self.setup.verbose >= 2:
                    print('Transformations will be fixed ' +
                          'during sampling.')
            if fix_m and not fix_t:
                method = sampler.sample_transformations_part
                if self.setup.verbose >= 2:
                    print('Mixture model parameters will be fixed ' +
                          'during sampling.')
            if fix_m and fix_t:
                raise ValueError(
                    'Cannot fix both mixture and transformations.')
            samplers = [method] * num_steps
        else:
            if self.setup.verbose >= 2:
                print('Global rotation sampling ' +
                      'with fixed translations.')
            num_outer_steps = (num_steps / self.setup.num_sub_steps) + 1
            samplers = []
            for i in range(num_outer_steps):
                for j in range(num_mixture_steps):
                    samplers.append(sampler.sample_mixture_part)
                samplers.append(global_sampler.sample)
                for j in range(num_transformation_steps - 1):
                    samplers.append(sampler.sample_rotations_part)
            samplers = samplers[:num_steps]

        return samplers

    def save_projections(self, model, grid, filename):
        """Save projections of given model.

        Project mixture along all rotations to obtain 2D mixture.
        Evaluate them on given grid.
        """
        projections = model.project()
        evaluated = [m.evaluate_grid_sparse(grid, 5.0)
                     for m in projections]
        GridDataList(evaluated).save_mrc(filename)

    def run_single(self, output_base, prior, data, grid2D):
        """Single reconstruction."""

        # Create initial model.
        model = self.create_initial_model(prior)

        # Compute projections of initial model to check it.
        self.save_projections(
            model, grid2D,
            output_base + '_initial_projections.mrc')

        # Initialise main sampler.
        std_range = self.setup.std_range
        sampler = IsoTransformSampler(data, prior, std_range)

        # Initialise global rotation sampler if necessary.
        if self.setup.use_global:
            global_sampler = BruteForceRotationSampler(
                data, self.setup.num_global_rotations)
        else:
            global_sampler = None

        # Prepare sampling sequence.
        samplers = self.prepare_samplers(sampler, global_sampler)

        # For collecting stats.
        stats = []

        # Generate samples.
        if self.setup.verbose >= 1:
            print 'Starting reconstruction (%s).' % output_base
        num_steps = self.setup.num_steps

        time_start = time.time()
        time_all_start = time.time()
        min_delay = 1.0 # Minimum time between updates in seconds.
        last_step = 0

        for i in range(num_steps):
            # Update model.
            samplers[i](model)
            # Compute stats (log-posterior, etc.).
            log_likelihood = model.log_likelihood(data.as_list())
            log_prior = prior.log_likelihood(model)
            log_posterior = log_prior + log_likelihood
            std = np.sqrt(model.mixture.variance)
            stats.append(np.array([i+1, log_posterior, std]))

            # Show how much time left, and some stats.
            elapsed = time.time() - time_start
            if (elapsed > min_delay and
                i % self.setup.num_sub_steps == 0):
                elapsed_total = time.time() - time_all_start
                time_per_step = elapsed_total / float(i+1)
                time_left = (num_steps - i - 1) * time_per_step
                if self.setup.verbose:
                    print(('Step %d/%d (%s): std = %1.4f, ' +
                           'log-posterior = %1.0f, ' +
                           'time remaining = %s') %
                          (i+1, num_steps, 
                           output_base,
                           std,
                           log_posterior, format_time(time_left)))
                time_start = time.time()
        total_time = time.time() - time_all_start
        if self.setup.verbose:
            print('Reconstruction time (%s) = %s.' % (
                output_base, format_time(total_time)))

        # Save stats.
        with open(output_base + '_stats.csv', 'wb') as fout:
            writer = csv.writer(fout)
            writer.writerow(['step', 'log-posterior', 'std'])
            writer.writerows(stats)

        # Save the final mixture and transformation parameters.
        model.mixture.save_csv(output_base + '_mixture.csv')
        Transformations(model.transformations).save_csv(
            output_base + '_transformations.csv')

        # Save the final model as mrc.
        r = (self.setup.input_grid_spacing /
             self.setup.output_grid_spacing)
        grid3D = Grids.create_centered(
            3, grid2D.shape[0] * r, grid2D.spacing[0] / r)
        gd = model.mixture.evaluate_grid_sparse(grid3D, std_range)
        gd.save_mrc(self.setup.output_mrc_filename)
        # Save mirror volume as well.
        gdm = GridDataReal(gd.grid, gd.reshaped()[::-1].flatten())
        gdm.save_mrc(output_base + '_mirror.mrc')

        # Compute projections of final model.
        self.save_projections(
            model, grid2D,
            output_base + '_final_projections.mrc')

    def run(self):
        """Run reconstruction."""

        # Load and prepare data (resize if necessary).
        # First as list of GridDataReal objects.
        # Then load 'num_counts' and convert to GridDataNestedList.
        gds = GridDataList.load_mrc(self.setup.input_filename,
                                    self.setup.input_grid_spacing)
        # Clip negative values to zero.
        gds.clip()
        if self.setup.verbose >= 2:
            grid_original = gds.get(0, 1.0).grid
            s = grid_original.shape[0]
            dx = grid_original.spacing[0]
            print(('Original input images have size %dx%d and ' +
                   'pixel size %1.3f Angstrom.') % (
                       s, s, dx))
        
        if self.setup.input_size:
            new_shape = np.ones(2, dtype='int') * self.setup.input_size
            if self.setup.verbose >= 2:
                print('Resizing input images to %dx%d.' % (
                    self.setup.input_size, self.setup.input_size))
            gds = GridDataList([gd.resize_centered(new_shape)
                                for gd in gds.as_list()])
        else:
            if self.setup.verbose >= 2:
                s = gds.get(0, 1.0).grid.shape[0]
                print('Using original input image sizes ' +
                      'of %dx%d.' % (s, s))
        gds.clip()

        data = gds.to_nested_list(self.setup.num_counts)
        grid2D = data.get(0, 1.0).grid
        if self.setup.verbose >= 2:
            s = grid2D.shape[0]
            dx = grid2D.spacing[0]
            print(('Running algorithm with images of size %dx%d and ' +
                   'pixel size %1.3f Angstrom.') % (s, s, dx))
        data.save_mrc(self.setup.output_base + '_input.mrc')

        # Construct prior.
        prior = self.construct_prior(
            data.get_num_projections(), grid2D)

        # Do reconstruction, either single or in parallel.
        if self.setup.num_parallel == 1:
            output_base = self.setup.output_base
            self.run_single(output_base, prior, data, grid2D)
        else:
            get_output_base = lambda i : (
                self.setup.output_base + '_%d' % i)

            def worker(index):
                np.random.seed()
                time.sleep(1.0 * index)
                output_base = get_output_base(index)
                self.run_single(output_base, prior, data, grid2D)

            jobs = []
            for i in range(self.setup.num_parallel):
                p = mp.Process(target=worker, args=(i, ))
                jobs.append(p)
                p.start()
            for p in jobs:
                p.join()

            # Choose model with highest log-posterior.
            def last_lp(i):
            # for i in range(self.setup.num_parallel):
                output_base = get_output_base(i)
                stats_dict = read_csv_dict(output_base + '_stats.csv')
                # print stats_dict.keys()
                return float(np.array(stats_dict['log-posterior'])[-1])
            lps = np.array(map(last_lp, range(self.setup.num_parallel)))
            max_index = np.argmax(lps)
            if self.setup.verbose >= 2:
                print ('Selecting reconstruction %d ' +
                       'as having the highest log-posterior.') % (
                           max_index)

            # Copy files.
            output_base_i = get_output_base(max_index)
            output_base = self.setup.output_base
            names = ['_initial_projections.mrc',
                     '_final_projections.mrc',
                     '_mirror.mrc',
                     '_mixture.csv',
                     '_stats.csv',
                     '_transformations.csv']
            for name in names:
                shutil.copy(output_base_i + name,
                            output_base + name)



