#
#    MiMiCPy: Python Based Tools for MiMiC
#    Copyright (C) 2020-2021 Bharath Raghavan,
#                            Florian Schackert
#
#    This file is part of MiMiCPy.
#
#    MiMiCPy is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Lesser General Public License as
#    published by the Free Software Foundation, either version 3 of
#    the License, or (at your option) any later version.
#
#    MiMiCPy is distributed in the hope that it will be useful, but
#    WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Lesser General Public License for more details.
#
#    You should have received a copy of the GNU Lesser General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

"""Module for gro files"""

import numpy as np
import pandas as pd
from ..utils.errors import ParserError
from .base import BaseCoordsClass

class Gro(BaseCoordsClass):
    """reads gro files"""

    def _read(self):
        """Read atom coordinates and box dimensions"""

        def mapped(value):
            if value.isnumeric():
                return np.nan
            try:
                return float(value)
            except ValueError:
                return np.nan

        def string_to_array(string):
            return np.array(list(map(mapped, string.split())))

        self.file.readline()
        number_of_atoms = int(self.file.readline())
        first_atom_line = self.file.readline()
        self.file.buffer = len(first_atom_line) * self.buffer

        number_of_rows = len(first_atom_line.split()) - 3
        if number_of_rows == 3:
            cols = ['x', 'y', 'z']
        elif number_of_rows == 6:
            cols = ['x', 'y', 'z', 'v_x', 'v_y', 'v_z']
        else:
            raise ParserError(self.file_name, details='Gro file is not formatted properly.')

        values = string_to_array(first_atom_line)
        for string in self.file:
            values = np.append(values, string_to_array(string))
        values = values[~np.isnan(values)]

        expected_len = number_of_atoms * number_of_rows
        if len(values) == expected_len:
            coords = values
            box = self.file.readline().split()
        elif len(values) == expected_len + 3:
            coords = values[:-3]
            box = values[-3:]
        else:
            raise ParserError(self.file_name, details='Gro file is not formatted properly.')

        coords = pd.DataFrame(coords.reshape(number_of_atoms, number_of_rows), columns=cols)
        coords['id'] = coords.index.to_numpy()+1

        return coords.set_index(['id']), box.tolist()

    def _write(self, mpt_coords, box, title):
        int_checker = lambda i: self.int_checker(i, 5)
        str_checker = lambda i: self.str_checker(i, 5)
        
        if title:
            title = title + ', '
        
        gro_str = '{}Generated by MiMiCPy\n{}\n'.format(title, len(mpt_coords))
        for i, r in mpt_coords.iterrows():
            gro_str += '{:5d}{:<5}{:>5}{:5d}{:8.3f}{:8.3f}{:8.3f}\n'.format(int_checker(r['resid']),\
                                                                     str_checker(r['resname']),\
                                                                     str_checker(r['name']),\
                                                                     int_checker(r['id']),\
                                                                 r['x'], r['y'], r['z'])

        if box is None:
            box = [0, 0, 0]
            for i, c in enumerate(['x', 'y', 'z']):
                box[i] = abs(max(mpt_coords[c]) - min(mpt_coords[c])) # find box size

        gro_str += '   {:.5f}   {:.5f}   {:.5f}\n'.format(box[0], box[1], box[2])

        return gro_str