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

A script for displaying and saving density distribution generated by
scft_confined/SlabXX3d, where XX stands for the polymer model. For A-B diblock
copolymer, XX=AB; For A-B diblock copolymer + grafted C polymer, XX=ABgC.

Copyright (C) 2014 <NAME> (<EMAIL>)

'''

import os
import argparse
import matplotlib.pyplot as plt
import numpy as np
from scipy.io import loadmat, savemat

import mpltex.acs
from chebpy import cheb_quadrature_clencurt
from scftpy import SCFTConfig
from scftpy import contourf_slab2d, list_datafile

parser = argparse.ArgumentParser(description='aslab2d options')

parser.add_argument('-b', '--batch',
                    action='store_true',
                    help='If present or True, perform batch mode.')
parser.add_argument('-r', '--render',
                    action='store_true',
                    help='If present or True, render density fields.')
parser.add_argument('-m', '--model',
                    default='AB',
                    help='scft model: AB, ABgC.')
parser.add_argument('-p', '--path',
                    default='.',
                    help='Path to be processed, for batch mode use.')
parser.add_argument('-d', '--data_file',
                    default='scft_out',
                    help='SCFT generated data file, *.mat')
args = parser.parse_args()


def plot(path, data):
    skipfiles = []
    datafiles = list_datafile(path, data)
    datafiles.sort()
    for f in skipfiles:
        for df in datafiles:
            p = os.path.dirname(df)
            d = os.path.basename(p)
            if d == f:
                datafiles.remove(df)
    Fs = []
    labels = []
    labels_tex = []

    fig1 = plt.figure()
    ax1 = fig1.add_subplot(111)  # F vs t
    fig2 = plt.figure()
    ax2 = fig2.add_subplot(111)  # F vs time
    fig3 = plt.figure()
    ax3 = fig3.add_subplot(111)  # err_res vs t
    fig4 = plt.figure()
    ax4 = fig4.add_subplot(111)  # err_res vs time
    fig5 = plt.figure()
    ax5 = fig5.add_subplot(111)  # err_phi vs t
    fig6 = plt.figure()
    ax6 = fig6.add_subplot(111)  # err_phi vs time
    fig7 = plt.figure()
    ax7 = fig7.add_subplot(111)  # err_F vs t
    fig8 = plt.figure()
    ax8 = fig8.add_subplot(111)  # err_F vs time

    for dfile in datafiles:
        mat = loadmat(dfile)
        p = os.path.dirname(dfile)
        label = os.path.basename(p)
        print label
        t, time, F = mat['t'], mat['time'], mat['F']
        t = t.reshape(t.size)
        time = time.reshape(time.size)
        F = F.reshape(F.size)
        err_res = mat['err_residual']
        err_res = err_res.reshape(err_res.size)
        try:
            err_phi = mat['err_phi']
        except:
            err_phi = np.ones_like(err_res)
        err_phi = err_phi.reshape(err_phi.size)
        for i in xrange(t.size):
            if t[i] == 0:
                imax = i
                break
        t = t[:imax]
        time = time[:imax]
        F = F[:imax]
        err_res = err_res[:imax]
        err_phi = err_phi[:imax]
        err_F = F[1:] - F[:-1]
        Fs.append(F[-1])
        labels.append(label)
        phiA_meanx = np.mean(mat['phiA'], axis=0)
        phiA_meanxy = np.mean(phiA_meanx, axis=0)
        phiA_mean = 0.5 * cheb_quadrature_clencurt(phiA_meanxy)
        print '\t', str(F[-1]), '\t', str(t[-1]), '\t', str(time[-1]),
        print '\t', str(phiA_mean)
        print '\t', str(err_res[-1])
        label = '$' + label + '$'
        labels_tex.append(label)
        ax1.plot(t, F, label=label)
        ax2.plot(time, F, label=label)
        ax3.plot(t, err_res, label=label)
        ax4.plot(time, err_res, label=label)
        ax5.plot(t, err_phi, label=label)
        ax6.plot(time, err_phi, label=label)
        ax7.plot(t[1:], np.abs(err_F), label=label)
        ax8.plot(time[1:], np.abs(err_F), label=label)

    datafile = os.path.join(path, 'data.mat')
    savemat(datafile, {'labels': labels, 'F': Fs})

    ax3.set_yscale('log')
    ax4.set_yscale('log')
    ax5.set_yscale('log')
    ax6.set_yscale('log')
    ax7.set_xscale('log')
    ax7.set_yscale('log')
    ax8.set_xscale('log')
    ax8.set_yscale('log')
    ax1.legend(loc='best')
    ax2.legend(loc='best')
    ax3.legend(loc='best')
    ax4.legend(loc='best')
    ax5.legend(loc='best')
    ax6.legend(loc='best')
    ax7.legend(loc='best')
    ax8.legend(loc='best')

    figfile = os.path.join(path, 'F-t.eps')
    fig1.savefig(figfile, format='eps', bbox_inches='tight')
    figfile = os.path.join(path, 'F-time.eps')
    fig2.savefig(figfile, format='eps', bbox_inches='tight')
    figfile = os.path.join(path, 'err_res-t.eps')
    fig3.savefig(figfile, format='eps', bbox_inches='tight')
    figfile = os.path.join(path, 'err_res-time.eps')
    fig4.savefig(figfile, format='eps', bbox_inches='tight')
    figfile = os.path.join(path, 'err_phi-t.eps')
    fig5.savefig(figfile, format='eps', bbox_inches='tight')
    figfile = os.path.join(path, 'err_phi-time.eps')
    fig6.savefig(figfile, format='eps', bbox_inches='tight')
    figfile = os.path.join(path, 'err_F-t.eps')
    fig7.savefig(figfile, format='eps', bbox_inches='tight')
    figfile = os.path.join(path, 'err_F-time.eps')
    fig8.savefig(figfile, format='eps', bbox_inches='tight')
    plt.close('all')

    x = np.arange(len(datafiles))
    Fs = np.array(Fs)
    fig9 = plt.figure(figsize=[10, 3])
    ax9 = fig9.add_subplot(111)  # F vs test
    ax9.plot(x, Fs, 'o')
    #plt.xscale('log')
    plt.xticks(x, labels_tex)
    figfile = os.path.join(path, 'F-all.eps')
    fig9.savefig(figfile, format='eps', bbox_inches='tight')
    plt.close()


def render(model, path, data):
    datafiles = list_datafile(path, data)
    for dfile in datafiles:
        p = os.path.dirname(dfile)
        #pfile = os.path.join(p, 'param.ini')
        #config = SCFTConfig.from_file(pfile)
        #model = config.model.model
        #Lx, Ly = config.grid.Lx, config.grid.Ly
        #La, Lb, Lc = config.uc.a, config.uc.b, config.uc.c

        mat = loadmat(dfile)
        if model == 'AB':
            phiA, phiB = mat['phiA'], mat['phiB']
        elif model == 'ABgC':
            phiA, phiB, phiC = mat['phiA'], mat['phiB'], mat['phiC']
        else:
            raise ValueError('Unknown model.')
        Lx, Ly, Lz = phiA.shape
        La, Lb, Lc = mat['a_curr'][0,0], mat['b_curr'][0,0], mat['c_curr'][0,0]

        yp, zp, phiApx = contourf_slab2d(phiA[Lx/2], Lb, Lc)
        figfile = os.path.join(p, 'phiAx.eps')
        plt.savefig(figfile, format='eps', bbox_inches='tight')
        plt.close()
        xp, zp, phiApy = contourf_slab2d(phiA[:, Ly/2, :], La, Lc)
        figfile = os.path.join(p, 'phiAy.eps')
        plt.savefig(figfile, format='eps', bbox_inches='tight')
        plt.close()
        yp, zp, phiBpx = contourf_slab2d(phiB[Lx/2], Lb, Lc)
        figfile = os.path.join(p, 'phiBx.eps')
        plt.savefig(figfile, format='eps', bbox_inches='tight')
        plt.close()
        xp, zp, phiBpy = contourf_slab2d(phiB[:, Ly/2, :], La, Lc)
        figfile = os.path.join(p, 'phiBy.eps')
        plt.savefig(figfile, format='eps', bbox_inches='tight')
        plt.close()
        if model == 'ABgC':
            yp, zp, phiCpx = contourf_slab2d(phiC[Lx/2], Lb, Lc)
            figfile = os.path.join(p, 'phiCx.eps')
            plt.savefig(figfile, format='eps', bbox_inches='tight')
            plt.close()
            xp, zp, phiCpy = contourf_slab2d(phiC[:, Ly/2, :], La, Lc)
            figfile = os.path.join(p, 'phiCy.eps')
            plt.savefig(figfile, format='eps', bbox_inches='tight')
            plt.close()

        plt.plot(zp, phiApx[Lx/2], label='$\phi_Ax$')
        plt.plot(zp, phiApy[Ly/2], label='$\phi_Ay$')
        plt.plot(zp, phiBpx[Lx/2], label='$\phi_Bx$')
        plt.plot(zp, phiBpy[Ly/2], label='$\phi_By$')
        if model == 'ABgC':
            plt.plot(zp, phiCpx[Lx/2], label='$\phi_Cx$')
            plt.plot(zp, phiCpy[Ly/2], label='$\phi_Cy$')
        plt.legend(loc='best')
        plt.xlabel('$z$')
        plt.ylabel('$\phi(z)$')
        figfile = os.path.join(p, 'profile.eps')
        plt.savefig(figfile, format='eps', bbox_inches='tight')
        plt.close()


if __name__ == '__main__':
    plot(args.path, args.data_file)
    if args.render:
        render(args.model, args.path, args.data_file)
