#!/usr/bin/env python3
from __future__ import print_function
from __future__ import division
import argparse
import pickle
import time
import CCBuilder as ccb
import CCBuilder_c as ccb_c
import numpy as np
import scipy.special


def uniform_dist(x):
    """ Returns uniform distributions of given range """
    return lambda: np.random.uniform(x[0], x[1])


def weibull_dist(a, mu):
    """ Returns Weibull distributions for given shape parameter and average """
    return lambda: np.random.weibull(a) * mu / scipy.special.gamma(1/a + 1)


def parse_dist(arg):
    # Parses input string for given distribution.
    # Returns a distribution, and the average
    d, params = arg.split(':')
    params = [float(x) for x in params.split(',')]
    if d == 'U':
        return uniform_dist(params), np.mean(params)
    elif d == 'W':
        a, mu = params
        return weibull_dist(a, mu), mu


parser = argparse.ArgumentParser(description='''Generate a WC microstructure.
Grain shape/size supports 2 types of distributions:
Uniform: U:low,high
Weibull: U:a,mu     (a=k in some notation, mu=mean)
''',
                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)

# parser.add_argument('-V', dest='verbose', action='store_true', help='Verbose mode.')

parser.add_argument('-f', dest='fname', metavar='basename', required=True, help='Output base filename.')

parser.add_argument('-L', dest='L', metavar='length', required=True, type=float, help='Cell length (volume is L^3)')

parser.add_argument('-m', dest='m', metavar='m', required=True, type=int,
                    help='Grid resolution. Total number of voxels are (m*L)^3')

parser.add_argument('--vol_frac_goal', dest="vol_frac_goal", metavar='v', required=True, type=float,
                    help='Goal for volume fraction WC (excluding overlap)')

parser.add_argument('-s', dest='seed', metavar='s', default=None, type=int,
                    help='Seed for RNG. Given identical parameters, ' +
                         'CCBuilder will generate identical output given a controlled seed.')

parser.add_argument('--stray_cleanup', action='store_true', help='Clean up stray voxels')

group = parser.add_argument_group('WC grain shape')
group.add_argument('-k', dest='k_dist', metavar='type,[params]', default='U:0.4,1.4',
                   help='k distribution')
group.add_argument('-r', dest='r_dist', metavar='type,[params]', default='U:0.1,0.4',
                   help='r distribution')
group.add_argument('-d', dest='d_dist', metavar='type,[params]', default='U:0.5,1.5',
                   help='d distribution')

group = parser.add_argument_group('Packing')
group.add_argument('--use_potential', action='store_true', help='Use repulsive potential.')
group.add_argument('--nr_tries', dest='nr_tries', metavar='n', default=2500, type=int,
                   help='Number of random translations.')
group.add_argument('--delta', dest='delta', metavar='d', type=float,
                   help='Maximum distance for randomized translations.')
group.add_argument('--m_coarse', dest="m_coarse", metavar='mc', default=10,
                   help='Grid resolution during packing.')

group = parser.add_argument_group('Potts simulation')
group.add_argument('--mc_steps', dest="mc_steps", metavar='steps', default=0.05, type=float,
                   help='Monte-Carlo steps (scales with (m*L)^4. Set to zero to turn off.')
group.add_argument('--tau', dest='tau', metavar='t', default=0.5, type=float,
                   help='Ficticious temperature in Potts model.')

options = parser.parse_args()

if options.seed is not None:
    np.random.seed(options.seed)

# Heuristic mapping from actual to goal volume fraction
# vol_frac_goal = (alpha - 2)/(2 * alpha) + 1/alpha * np.sqrt(1 - alpha * np.log(-2*(vol_frac - 1)))

d_eq, d_0 = parse_dist(options.d_dist)
r, r_0 = parse_dist(options.r_dist)
k, k_0 = parse_dist(options.k_dist)

fname = options.fname

# to avoid confusion with types:
m = np.int(options.m)
m_coarse = np.int(options.m_coarse)
L = np.float(options.L)
mc_steps = np.float(options.mc_steps)
vol_frac_goal = np.double(options.vol_frac_goal)
tau = np.double(options.tau)
nr_tries = np.int(options.nr_tries)
delta_x = d_0/float(m)
M = np.int(m * L / d_0)
M_coarse = np.int(m_coarse * L / d_0)

idelta = M
idelta_coarse = M_coarse
if options.delta:
    idelta = np.int(M * options.delta / L)
    idelta_coarse = np.int(M_coarse * options.delta / L)

trunc_triangles = ccb.prepare_triangles(vol_frac_goal, L, r, k, d_eq)
# trunc_triangles = trunc_triangles[:1]
# trunc_triangles[0].rot_matrix = np.eye(3)
# trunc_triangles[0].rot_matrix_tr = np.eye(3)
# trunc_triangles[0].midpoint = np.array([2., 2., 2.])

# Sort triangles w.r.t. volume, so that large triangles are added to the box first (better packing)
trunc_triangles.sort(key=lambda x: x.volume, reverse=True)
print('Prepared', len(trunc_triangles), 'triangles')

if options.use_potential:
    ccb.optimize_midpoints(L, trunc_triangles)

if m_coarse == m:
    grain_ids, overlaps, voxel_indices = ccb_c.populate_voxels(M, L, trunc_triangles, nr_tries, idelta, 1.0)
else:
    if nr_tries > 0:
        # Optimization: Use coarser grid for packing, then insert packed grains into fine grid
        # No need to get the return values, trunc_triangles
        ccb_c.populate_voxels(M_coarse, L, trunc_triangles, nr_tries, idelta_coarse, 1.0)
    grain_ids, overlaps, voxel_indices = ccb_c.populate_voxels(M, L, trunc_triangles, 1, 0, 1.0)

if mc_steps > 0:
    start_time = time.time()
    # Do Potts on coarse grid first for an improved initial guess.
    M_coarseMC = M//2
    grain_ids_coarse, overlaps_coarse, voxel_indices_coarse = ccb_c.populate_voxels(M_coarseMC, L, trunc_triangles, 0, 0, 1.0)
    _, gb_voxels_coarse, _ = ccb_c.calc_surface_prop(M_coarseMC, grain_ids_coarse)
    ccb_c.make_mcp_bound(M_coarseMC, grain_ids_coarse, gb_voxels_coarse, overlaps_coarse, voxel_indices_coarse,
                         np.int(mc_steps * M_coarseMC**4), tau)

    # Copy over that solution to the overlap regions of the fine grid as a starting point
    M2 = M**2
    i = np.nonzero(overlaps)[0]
    iz = i // M2
    iy = (i - iz*M2) // M
    ix = i - iz*M2 - iy*M
    cix = ix * M_coarseMC // M
    ciy = iy * M_coarseMC // M
    ciz = iz * M_coarseMC // M
    ci = cix + ciy*M_coarseMC + ciz*M_coarseMC**2
    gid = grain_ids_coarse[ci]
    # Could use a Cython implementation for efficiency.
    for ii, g in zip(i, gid):
        if g != grain_ids[ii] and np.searchsorted(voxel_indices[g-2], ii) < len(voxel_indices[g-2]):
            grain_ids[ii] = g
    # This might change a few voxels to a value that they shouldn't obtain, but it's barely noticeable
    # grain_ids_1[i] = grain_ids_coarse[ci]
    _, gb_voxels, _ = ccb_c.calc_surface_prop(M, grain_ids)
    # and run the full resolution MCP:
    ccb_c.make_mcp_bound(M, grain_ids, gb_voxels, overlaps, voxel_indices, np.int(mc_steps * M ** 4), tau)
    print('Potts model took {} seconds'.format(np.str(time.time() - start_time)))

if options.stray_cleanup:
    start_time = time.time()
    ccb_c.stray_cleanup(M, grain_ids)
    print('Stray voxel cleanup took {} seconds'.format(np.str(time.time() - start_time)))

surface_voxels, gb_voxels, interface_voxels = ccb_c.calc_surface_prop(M, grain_ids)
phases, good_voxels, euler_angles = ccb_c.calc_grain_prop(M, grain_ids, trunc_triangles)

phase_volumes = np.bincount(phases)
vol_frac_WC = phase_volumes[2] / np.float(M ** 3)
vol_frac_Co = 1 - vol_frac_WC
mass_frac_WC = ccb.mass_fraction(vol_frac_WC)

sum_gb_voxels = np.sum(gb_voxels)
contiguity = sum_gb_voxels / np.float(sum_gb_voxels + np.sum(interface_voxels))

print('Contiguity {:5f}, Co volume frac {:.5f}, mass frac {:.5f}'.format(
    contiguity, 1 - vol_frac_WC, ccb.mass_fraction(vol_frac_WC)))

ccb.write_dream3d(fname, 3 * [M], 3 * [delta_x], trunc_triangles, grain_ids, phases, good_voxels,
                  euler_angles, surface_voxels, gb_voxels, interface_voxels, overlaps)

with open(fname + '_trunc_triangles.data', 'wb') as f:
    pickle.dump([t.rot_matrix for t in trunc_triangles], f)

# Saving grain volume data
if False:
    grain_volumes = np.bincount(grain_ids)
    d_eq = ccb.volume_to_eq_d(grain_volumes[2:] * delta_x ** 3)
    # np.savetxt(fname + '_d_orig.txt', [t.d_eq for t in trunc_triangles])
    np.savetxt(fname + '_d.txt', d_eq)
    # Plot initial and final distributions
    import matplotlib.pyplot as plt
    plt.hist(np.array([t.d_eq for t in trunc_triangles]), alpha=0.5, bins=15, normed=True, label='Initial')
    plt.hist(d_eq, alpha=0.5, bins=15, normed=True, label='Final')
    plt.legend(loc='upper right')
    plt.show()
