# Copyright (C) 2010 by CAMd, DTU
# Please see the accompanying LICENSE file for further information.

# This file is taken (almost) verbatim from CMR with D. Landis agreement

FIELD_SEPARATOR = "\\"
PARA_START = "\n\n"
PARA_END = "\\\\@"

names = ['', '', 'Computer_system', 'Type_of_run', 'Method', 'Basis_set',
         'Chemical_formula', 'Person', 'Date', '', '', '', '', 'Title', '']
names_compact = ['', '', 'Computer_system', 'Type_of_run', 'Method',
                 'Basis_set', 'Chemical_formula', 'Person', 'Date', '', '', '',
                 '', 'Title', '']

charge_multiplicity = 15


class GaussianReader:

    def auto_type(self, data):
        """ tries to determine type"""
        try:
            return float(data)
        except ValueError:
            pass

        try:
            ds = data.split(",")
            array = []

            for d in ds:
                array.append(float(d))

            return array
        except ValueError:
            pass

        return data

    def __init__(self, filename, read_structures=False):
        """filename is NOT optional"""
        if isinstance(filename, str):
            fileobj = open(filename, 'r')
        elif hasattr(filename, 'seek'):
            fileobj = filename
            fileobj.seek(0)  # Re-wind fileobj
        else:
            raise RuntimeError('Cannot use given filename, make sure it is a '
                               'string or a fileobject')

        content = fileobj.read()

# handles the case that users used windows after the calculation:
        content = content.replace("\r\n", "\n")

        self.parse(content)

        # read structures from file
        if read_structures:
            self.read_structures(content)

    def get_structures(self, content=None):
        """Get Structures"""
        if hasattr(self, 'structures'):
            return self.structures
        elif content is None:
            raise RuntimeError('Images not available and no content parsed!')
        else:
            self.read_structures(content)
            return self.structures

    def read_structures(self, content=None):
        """Read Structures from file and wirte them to self.structures"""
        from ase.atoms import Atoms
        from ase.atom import Atom
        images = []
        temp_items = content.split('Standard orientation')[1:]
        for item_i in temp_items:
            lines = [line for line in item_i.split('\n') if len(line) > 0]
            # first 5 lines are headers
            del lines[:5]
            images.append(Atoms())
            for line in lines:
                # if only - in line it is the end
                if set(line).issubset(set('- ')):
                    break
                tmp_line = line.strip().split()
                if not len(tmp_line) == 6:
                    raise RuntimeError('Length of line does not match '
                                       'structure!')

                # read atom
                try:
                    atN = int(tmp_line[1])
                    pos = tuple(float(x) for x in tmp_line[3:])
                except ValueError as e:
                    raise ValueError('Expected a line with three integers and '
                                     'three floats.') from e
                images[-1].append(Atom(atN, pos))
        self.structures = images
        return

    def parse(self, content):
        from ase.data import atomic_numbers
        chg_mult = charge_multiplicity
        self.data = []
        temp_items = content.split(PARA_START)
        seq_count = 0
        for i in temp_items:
            i = i.replace("\n ", "")
            if i.endswith(PARA_END):
                i = i.replace(PARA_END, "")
                i = i.split(FIELD_SEPARATOR)

                new_dict = {}
                self.data.append(new_dict)

                new_dict['Sequence number'] = seq_count
                seq_count += 1
                for pos in range(len(names)):
                    if names[pos] != "":
                        # hack, since this section is too short
                        # if there is no title
                        if names[pos] == "Title" and i[pos] == "":
                            chg_mult -= 1
                            break
                        new_dict[names[pos]] = self.auto_type(i[pos])

                chm = i[chg_mult].split(",")
                new_dict["Charge"] = int(chm[0])
                new_dict["Multiplicity"] = int(chm[1])

# Read atoms
                atoms = []
                positions = []
                position = chg_mult + 1
                while position < len(i) and i[position] != "":
                    s = i[position].split(",")
                    atoms.append(atomic_numbers[s[0].capitalize()])
                    # if fragments are specified, there are 4 numbers
                    # first one integer and then xyz coords
                    # therefore use xyz from the end
                    positions.append([float(s[-3]), float(s[-2]), float(s[-1])])
                    position = position + 1

                new_dict["Atomic_numbers"] = atoms
                new_dict["Positions"] = positions
# Read more variables
                position += 1
                while position < len(i) and i[position] != "":
                    s = i[position].split('=')
                    if len(s) == 2:
                        new_dict[s[0]] = self.auto_type(s[1])
                    else:
                        print("Warning: unexpected input ", s)
                    position = position + 1

    def __iter__(self):
        """returns an iterator that iterates over all keywords"""
        return self.data.__iter__()

    def __len__(self):
        return len(self.data)

    def __getitem__(self, pos):
        return self.data[pos]