#!/usr/bin/env python
"""
Calculate elastic constants C_ij by least square fitting
using Green-Lagrange deformation tensors, which is the same
scheme as materialsproject.org does.

PREPARE mode creates directories and POSCAR files to be computed.
ANALYZE mode reads the stress values obtained by some code.

Usage:
    elasticity.py prepare [options] POSCAR
    elasticity.py analyze [options] POSCAR STRSFILE

Options:
  -h, --help  Shows this message and exit.
  --delta1 DELTA1
              Maximum strain value of diagonal elements. Available only with prepare. [default: 0.002]
  --delta2 DELTA2
              Maximum strain value of off-diagonal elements. Available only with prepare. [default: 0.01]
"""
from __future__ import print_function

import os
import numpy as np
from docopt import docopt
from scipy.optimize import curve_fit
import spglib
import yaml

import ase.io
from nappy.napsys import NAPSystem
from nappy.io import read, from_ase
import copy

__author__  = 'Ryo KOBAYASHI'
__version__ = '200521'
__licence__ = 'MIT'

_confname = 'conf.elast.yaml'

#...constants
_prefix = 'elast_'

def quad_func(x,a,b):
    return a *x**2 +b

def get_erg(fname):
    with open(fname) as f:
        l = f.readline()
        erg = float(l.split()[0])
    return erg

def get_deformations(dlt1max,dlt2max):

    dlt1s = [-dlt1max, -dlt1max/2, dlt1max/2, dlt1max]
    dlt2s = [-dlt2max, -dlt2max/2, dlt2max/2, dlt2max]

    fmats = []
    for dlt1 in dlt1s:
        for t in range(3):
            fmat = np.identity(3,dtype=float)
            fmat[t,t] += dlt1
            fmats.append(fmat)
            
    #elems = ((1,2),(0,2),(0,1),(2,1),(2,0),(1,0))
    elems = ((1,2),(0,2),(0,1))
    for dlt2 in dlt2s:
        for e in elems:
            fmat = np.identity(3,dtype=float)
            fmat[e] = dlt2
            fmats.append(fmat)
    return fmats

def prepare(infname='POSCAR',dlt1max=0.01,dlt2max=0.06):

    #...original system
    #nsys0 = NAPSystem(fname=infname)
    nsys0 = read(infname)
    #orig_atoms = read(infname,format='vasp')
    orig_atoms = nsys0.to_ase_atoms()

    #...get deformations
    fmats = get_deformations(dlt1max,dlt2max)

    #...deform original system and save to _prefix_##/POSCAR
    for i,fmat in enumerate(fmats):
        atoms = orig_atoms.copy()
        #nsys = copy.deepcopy(nsys0)
        dname = _prefix +"{0:02d}".format(i)
        os.system('mkdir -p {}'.format(dname))
        print(dname)
        cell0 = atoms.get_cell()
        emat = 0.5 *(np.dot(fmat.T,fmat) -np.identity(3)) +np.identity(3)
        cell = np.dot(emat,cell0)
        #print(i,emat,cell)
        # cell = np.dot(cell0,fmat.T)
        atoms.set_cell(cell,scale_atoms=True)
        atoms.write(dname+'/POSCAR',format='vasp',vasp5=True,direct=True,
                    sort=False)
    
    print('prepare done')
    print('')
    print('After performing VASP or pmd calculations in these directories, '
          +'run the following command:')
    print('  $ python elasticity.py analyze str.ref')
    print('or')
    print('  $ python elasticity.py analyze strs.pmd')
    print('')

def cdote(strns,*params):
    """
    Compute C_ij*e_j to get s_i.
    """
    ctnsr = params2ctnsr(params)
    
    strss = np.zeros((len(strns),6),dtype=float)
    for i,strn in enumerate(strns):
        strs = np.dot(ctnsr,strn)
        strss[i] = strs
    return strss.flatten()

def func(x,*args):
    """
    Objective function to be minimized in least square fitting.

    .. math::

      L= \sum_s 1/N_s \sum_i 1/6 [ S_i^s -\sum_j C_{ij} e_j^s]^2

    
    Options
    -------
    x : N-dimensional array
        Variable array to be optimized.
    args : 2-dimensional array of parameter arrays