from Elasticipy.tensors.second_order import SecondOrderTensor, SymmetricSecondOrderTensor
from Elasticipy.tensors.stress_strain import StressTensor
import pandas as pd
import numpy as np

def _rebuild_tensor(F11, F22, F33, F12, F13, F21, F23, F31, F32):
    return np.array([[F11, F12, F13], [F21, F22, F23], [F31, F32, F33]]).transpose((2,0,1))

def from_quadrature_file(file, returns='stress'):
    """
    Read data from quadrature output file generated by PRISMS plasticity.

    These files, usually named QuadratureOutputsXXX.csv (where XXX denotes the time increment), contain large amount of
    data, such as gradient components, stress values, grain ID etc.

    Parameters
    ----------
    file : str
        Path to quadrature file
    returns : str or list of str or tuple of str
        name(s) of requested field(s). They can be:
          - grain ID
          - phase ID
          - det(J)
          - twin
          - coordinates
          - orientation
          - elastic gradient
          - plastic gradient
          - stress
    Returns
    -------
    ndarray or SecondOrderTensor or StressTensor
        The number of returned values depends on the requested fields.
    """
    data=pd.read_csv(file, header=None, dtype=float, usecols=range(0,37))
    grainID, phaseID, detJ, twin, x, y, z, rot1, rot2, rot3, *other = data.T.to_numpy()
    Fe11, Fe22, Fe33, Fe12, Fe13, Fe21, Fe23, Fe31, Fe32, *FpStress = other
    Fp11, Fp22, Fp33, Fp12, Fp13, Fp21, Fp23, Fp31, Fp32, *stress_compo = FpStress
    if isinstance(returns, str):
        returns_list = (returns,)
    else:
        returns_list = returns
    returned_values = []
    for r in returns_list:
        rl = r.lower()
        if (rl == 'grain id') or (rl == 'grainid'):
            returned_values.append(grainID)
        elif (rl == 'phase id') or (rl == 'phaseid'):
            returned_values.append(phaseID)
        elif rl == 'det(j)':
            returned_values.append(detJ)
        elif rl == 'twin':
            returned_values.append(twin)
        elif rl == 'coordinates':
            returned_values.append(np.array([x, y, z]).T)
        elif rl == 'orientation':
            returned_values.append(np.array([rot1, rot2, rot3]).T)
        elif rl == 'elastic gradient':
            Fe = _rebuild_tensor(Fe11, Fe22, Fe33, Fe12, Fe13, Fe21, Fe23, Fe31, Fe32)
            returned_values.append(SecondOrderTensor(Fe))
        elif rl == 'plastic gradient':
            Fp = _rebuild_tensor(Fp11, Fp22, Fp33, Fp12, Fp13, Fp21, Fp23, Fp31, Fp32)
            returned_values.append(SecondOrderTensor(Fp))
        elif rl == 'stress':
            stress = _rebuild_tensor(*stress_compo)
            returned_values.append(StressTensor(stress))
        else:
            raise ValueError('Unknown return type')
    if isinstance(returns, str):
        return returned_values[0]
    else:
        return tuple(returned_values)

def from_stressstrain_file(file):
    """
    Read data from stress-strain file generated by PRISMS plasticity.

    This file is usually named stressstrain.txt

    Parameters
    ----------
    file : str
        Path to stress-strain file

    Returns
    -------
    SymmetricSecondOrderTensor
        Average Green-Lagrange strain tensor
    StressTensor
        Average stress tensor
    """
    data = pd.read_csv(file, sep='\t')
    Exx, Exy, Exz = data['Exx'], data['Exy'], data['Exz']
    Eyy, Eyz = data['Eyy'], data['Eyz']
    Ezz = data['Ezz']
    E = np.array([[Exx, Exy, Exz], [Exy, Eyy, Eyz], [Exz, Eyz, Ezz]]).transpose((2,0,1))
    Txx, Txy, Txz = data['Txx'], data['Txy'], data['Txz']
    Tyy, Tyz = data['Tyy'], data['Tyz']
    Tzz = data['Tzz']
    T = np.array([[Txx, Txy, Txz], [Txy, Tyy, Tyz], [Txz, Tyz, Tzz]]).transpose((2, 0, 1))
    return SymmetricSecondOrderTensor(E), StressTensor(T)

