"""
This is used for reading the output from the calculation in a nympy format.

Can be used interactively in ipython, allows for easy analysis and plotting...

The main subroutine is read_result, which reads the results and stores it as 
a class result.

Examples:

    res = read_result()

    To get a tensor with gamma=0.001, formula number 2, projection 1 and nk=100**3 do:

        res.get_result((0.001,2,1,100**3))

    For a spin-Hall effect there are no projections so do:

        res.get_result((0.001,2,100**3))
"""
import numpy as np
import xarray as xr
import itertools
from collections import OrderedDict

def read_result(dire='.',fil='linres_out'):
    dire + '/' + fil
    with open(dire + '/' + fil) as f:
        lines = f.readlines()

    for i,line in enumerate(lines):
        if 'Version' in line:
            version = line.split(':')[1].strip()
            if len(version.split('.')) != 3:
                version = lines[i+1].strip()


    vs = version.split('.')
    for i,v in enumerate(vs):
        vs[i] = int(v)
    if version == '0.1.0':
        raise Exception('Use analyze_res_compat!')
    elif vs[0] == 0 and vs[1] == 2 and vs[2] == 0:
        raise Exception('Use analyze_res_compat!')
    elif vs[0] == 0 and vs[1] >= 2 and vs[1] <= 3:
        raise Exception('Use analyze_res_compat!')
    elif vs[0] == 0 and vs[1] >= 2 and vs[1] <= 4 and vs[2] <= 12:
        raise Exception('Use analyze_res_compat!')
    else:
        return read_new(dire=dire,fil=fil)

def read_new(content=None,dire='.',fil='linres_out',empty_projs=False):
    if content is not None:
        lines = content
    else:
        with open(dire + '/' + fil) as f:
            lines = f.readlines()

    params = OrderedDict()
    start = False
    start_pranges = False
    start_X = False
    popped = []
    for i,line in enumerate(lines):
        if 'Unformatted output' in line:
            start = True
            continue
        if start:
            if 'rank' in line:
                rank = int(lines[i+1])
            if 'Parameter ranges' in line:
                start_pranges = True
                continue
            if start_pranges:
                if ':' in line:
                    if 'Param:' in line:
                        start_pranges = False
                        start_X = True
                        if len(params['Projection']) == 1:
                            if not empty_projs:
                                for i,p in enumerate(params):
                                    if p == 'Projection':
                                        popped.append(i)
                                params.pop('Projection')
                        if 'Omega' in params:
                            if len(params['Omega']) == 1:
                                if not False in [f < 100 for f in params['Formula']]:
                                    for i,p in enumerate(params):
                                        if p == 'Omega':
                                            popped.append(i)
                                    params.pop('Omega')
                        n_params = len(params)
                        param_lens = [len(params[param]) for param in params]
                        param_lens += [3]*rank
                        X = np.zeros(param_lens)
                    else:
                        params[line.strip().rstrip(':')] = []
                else:
                    ls = line.strip()
                    try:
                        val = int(ls)
                    except:
                        val = float(ls)
                    params[next(reversed(params))].append(val)
            if start_X:
                if 'Param:' in line:
                    X_p = np.zeros(3**rank)
                    ind = line.strip().lstrip('Param:').split()
                    ind = [int(i)-1 for i in ind]
                    for i in popped:
                        ind.pop(i)
                    ind2 = 0
                else:
                    try:
                        X_p[ind2] = float(line.strip())
                    except:
                        break
                    if ind2 == 3**rank - 1: 
                        ind3 = tuple(ind)+(Ellipsis,)
                        X[ind3] = np.reshape(X_p,(3,)*rank,order='F')
                    ind2 += 1

    coords = list(params.values())+[range(3)]*rank
    dims = list(params.keys()) + ['dim'+str(i+1) for i in range(rank)]
    Xa = xr.DataArray(X,coords=coords,dims=dims)

    return Xa

def read_k_resolved(path):
 
    spec_path = path.replace('.out','.spec')
    with open(spec_path) as f:
        lines = f.readlines()

    mode = None
    for line in lines:
        if 'Mode:' in line:
            mode = line.split()[1]
        if 'Version' in line:
            version = line.split(':')[1].strip()
            if len(version.split('.')) != 3:
                version = lines[i+1].strip()
        if 'Shape:' in line:
            shape = line.split()[1:]
            shape = [int(x) for x in shape]
        if 'Tensor rank:' in line:
            rank = int(line.split(':')[1])
        if 'k-dimension:' in line:
            kdim = int(line.split(':')[1])

    if mode is None:
        mode = 'k_resolved'

    if mode == 'k_resolved':
        shape2 = list(shape)
        shape2[-2:-2] = [3]*rank
        shape2.pop(-2)

    if version == '0.4.15':
        X = np.fromfile(path,np.float64)
        Xr = np.reshape(X,shape2,order='F')
    else:
        X = np.fromfile(path,np.float64)
        if mode == 'k_resolved':
            Xr1 = np.reshape(X,shape)
            Xr = np.reshape(Xr1,shape2,order='F')
        elif mode == 'bands':
            Xr = np.reshape(X,shape)
        else:
            raise Exception('Unknown mode')

    ks =['k1','k2','k3']
    ks = ks[0:kdim]

    if mode == 'k_resolved':
        dims = ks + ['dim'+str(i+1) for i in range(rank)] + ['param']
        coords = []
        for i in range(len(shape2)):
            coords.append((dims[i],range(shape2[i])))
    elif mode == 'bands':
        dims = ks + ['orb']
        coords = []
        for i in range(len(shape)):
            coords.append((dims[i],range(shape[i])))

    Xa = xr.DataArray(Xr,dims=dims,coords=coords)
    return Xa
