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)

