"""
ndxParser module is part of DuIvyTools for parsing the ndx file generated by GROMACS.

Written by DuIvy and provided to you by GPLv3 license.
"""

import os
import sys
import math
from typing import List, Union

base = os.path.dirname(os.path.realpath(os.path.join(__file__, "..")))
if base not in sys.path:
    sys.path.insert(0, base)

from utils import log


class NDX(log):
    """NDX class was designed to parse index file"""

    def __init__(self, ndxfile: str, new_file: bool = False) -> None:
        self.ndxfile: str = ndxfile
        self.indexs: List[int] = []
        self.names: List[str] = []
        self.column_nums: List[int] = []

        if not new_file and ndxfile:
            if not os.path.exists(ndxfile):
                self.error(f"No {ndxfile} detected! check it!")
            if ndxfile[-4:] != ".ndx":
                self.error(
                    f"you must specify a file with suffix .ndx, instead of {ndxfile}"
                )
            with open(ndxfile, "r") as fo:
                lines = [line.strip() for line in fo.readlines()]
            for id, line in enumerate(lines):
                if line.strip() == "":
                    continue
                elif line[0] == "[" and line[-1] == "]":
                    name = line[1:-1].strip()
                    self.indexs.append([])
                    self.column_nums.append(0)
                    self.names.append(name)
                else:
                    try:
                        line_list = [int(i) for i in line.split()]
                        if len(self.indexs[-1]) == 0:
                            self.column_nums[-1] = len(line_list)
                        self.indexs[-1] += line_list
                    except:
                        self.error(f"Unable to parse line {id} of {ndxfile}, check it!")
        if len(self.names) != len(self.indexs) or len(self.names) != len(
            self.column_nums
        ):
            self.critical("wrong length in paring ndx file")

    def __len__(self) -> int:
        """return number of groups"""
        return len(self.names)

    def __getitem__(
        self, key: Union[str, int]
    ) -> Union[Union[str, None], Union[None, List[int]]]:
        """get item by group id or group name"""
        if isinstance(key, int):
            if key < len(self):
                name = self.names[key]
                index = self.indexs[key]
                return name, index
            else:  # index over range
                return None, None
        elif isinstance(key, str):
            if key in self.names:
                index = self.indexs[self.names.index(key)]
                return key, index
            else:
                return None, None
        else:
            return None, None

    def __setitem__(self, key: Union[str, int], indexs: List[int]) -> None:
        """set item by group id of group name"""
        if isinstance(key, int):
            if key >= len(self):
                self.error("key over range for setting item of NDX")
            self.names[id] = key
            self.indexs[id] = indexs
            self.column_nums[id] = 15
        elif isinstance(key, str):
            if key not in self.names:
                self.indexs.append(indexs)
                self.names.append(key)
                self.column_nums.append(15)
            else:
                ids = self.get_id_by_name()
                for id in ids:
                    self.names[id] = key
                    self.indexs[id] = indexs
                    self.column_nums[id] = 15
        else:
            pass

    def add(self, name: str, indexs: List[int], column_num: int = 15) -> None:
        """add a new group to NDX"""
        self.names.append(name)
        self.indexs.append(indexs)
        self.column_nums.append(column_num)

    def __delitem__(self, key: Union[str, int]) -> None:
        """delete group by group id or group name"""
        if isinstance(key, int):
            if key >= len(self):
                self.error("key over range for deleting item of NDX")
            del self.names[key]
            del self.indexs[key]
            del self.column_nums[key]
        elif isinstance(key, str):
            name = key
            while name in self.names:
                id = self.names.index(name)
                del self.names[id]
                del self.indexs[id]
                del self.column_nums[id]
        else:
            self.error("Wrong key for deleting item of NDX")

    def get_id_by_name(self, key: str) -> List[int]:
        """get group id by group name"""
        id_list = []
        for id, name in enumerate(self.names):
            if key == name:
                id_list.append(id)
        return id_list

    @property
    def show_names(self) -> str:
        """convert group names and infos into string and return"""
        output: str = ""
        for id, name in enumerate(self.names):
            output += f"{id} {name:<20} {len(self.indexs[id])}\n"
        return output

    def save(self, outfile) -> None:
        """dump NDX class into index file"""
        with open(outfile, "w") as fo:
            fo.write(str(self))

    def __str__(self) -> str:
        """convert NDX class into string"""
        output: str = ""
        for id in range(len(self)):
            out = self.formatter(id)
            output += out
        output += "\n "
        return output

    def formatter(self, id: int) -> str:
        """format group into string by group id"""
        name = self.names[id]
        index = self.indexs[id]
        column_num = self.column_nums[id]
        if column_num <= 0:
            self.error(f"Unable to format indexs by 0 column at group id {id}")
        output: str = f"[ {name} ] \n"
        count = math.ceil(len(index) / column_num)
        for c in range(count):
            items = [f"{v:>4d}" for v in index[c * column_num : (c + 1) * column_num]]
            output += " ".join(items) + "\n"
        return output
