from __future__ import annotations

import pathlib
import re
import warnings
from io import StringIO
from typing import Optional, Tuple

import numpy as np
from numpy import log10

from ..configs import config_paths
from ..lib.logger import logger


class GraceEditor:
    _RE_HEADER_START = re.compile(r'# Grace project file')
    _RE_HEADER_END = re.compile(r'@timestamp def')
    _RE_OBJECT_START = re.compile(r'@with (box|ellipse|string|)')
    _RE_REGION_START = re.compile(r'@r(\w+) (on|off)', re.IGNORECASE)
    _RE_GRAPH_START = re.compile(r'@g(\w+) (on|off)', re.IGNORECASE)
    _RE_SET_START = re.compile(r'@target G(\d+).S(\d+)', re.IGNORECASE)
    _RE_SET_END = re.compile(r'^\s*&\s*$')
    _RE_COLOR = re.compile(r'@map color (?P<id>\d+) to '
                           r'\((?P<red>\d+), (?P<green>\d+), (?P<blue>\d+)\),\s+\"(?P<disp>.+)\"',
                           re.MULTILINE)
    _RE_FONT = re.compile(r'@map font (?P<id>\d+) to \"(?P<name>.+?)\",\s+\"(?P<disp>.+)\"')
    _RE_PAGE_SIZE = re.compile(r'@page size (\d+), (\d+)')

    def __init__(self, filename=None):
        self.header = GraceHeader()
        self.drawing_objects = []
        self.regions = []
        self.graphs = []
        self.sets = []

        _default_file = config_paths().joinpath('Default.agr')

        if filename == _default_file:
            self.defaults = self
        else:
            self.defaults = GraceEditor(filename=_default_file)

        self.file = filename
        if filename is not None:
            self.parse(filename)

    def __call__(self, filename: str):
        self.clear()
        self.parse(filename)

        return self

    def __str__(self):
        s = str(self.header)
        s += ''.join(map(str, self.regions))
        s += ''.join(map(str, self.drawing_objects))
        s += ''.join(map(str, self.graphs))
        s += ''.join(map(str, self.sets))

        return s

    def clear(self):
        self.header = GraceHeader()
        self.drawing_objects = []
        self.regions = []
        self.graphs = []
        self.sets = []
        self.graphs = []

    def new_graph(self):
        if self.file is None:
            self.parse(config_paths().joinpath('Default.agr'))
        else:
            max_idx = 0
            for g in self.graphs:
                max_idx = max(g.idx, max_idx)
            self.graphs.append(self.defaults.graphs[-1].copy())
            self.graphs[-1].renumber(max_idx+1)

        return self.graphs[-1]

    def new_drawing(self, dtype: str = 'line'):
        obj = GraceDrawing(dtype)
        self.drawing_objects.append(obj)
        return self.drawing_objects[-1]


    def new_set(self, graph):
        s = None
        g_idx = -1
        for g in self.graphs:
            if g.idx == graph:
                s = g.new_set()
                g_idx = g.idx
                break

        if s is None:
            raise ValueError(f'Graph {graph} was not found.')

        self.sets.append(GraceSet(g_idx, s.idx))
        self.sets[-1].append(f'@target G{g_idx}.S{s.idx}\n')
        self.sets[-1].append('@type xy\n')
        self.sets[-1].append('&\n')

        return s

    def parse(self, filename: str | pathlib.Path):
        self.file = pathlib.Path(filename)

        # we start always with the header
        current_pos = 'header'
        with self.file.open('r') as f:
            for line_nr, line in enumerate(f):

                # lots of states to check
                # agr file order: header, drawing object, region, graph, set
                if current_pos == 'header':
                    self.header.append(line)

                    # already at end of header
                    if self._RE_HEADER_END.match(line):
                        current_pos = 'header_end'
                    elif self._RE_GRAPH_START.match(line):
                        current_pos = self._make_graph(line)

                elif current_pos == 'header_end':
                    # what comes after the header? region, graph or drawing object?
                    if self._RE_REGION_START.match(line):
                        current_pos = 'region'
                        self.regions.append(GraceRegion())
                        self.regions[-1].append(line)
                    elif self._RE_GRAPH_START.match(line):
                        current_pos = self._make_graph(line)
                    elif self._RE_OBJECT_START.match(line):
                        current_pos = 'drawing'
                        self.drawing_objects.append(GraceDrawing())
                        self.drawing_objects[-1].append(line)
                    else:
                        logger.warn('What happened with this line?', line.rstrip())

                elif current_pos == 'drawing':
                    if self._RE_REGION_START.match(line):
                        current_pos = 'region'
                        self.regions.append(GraceRegion())
                        self.regions[-1].append(line)
                    else:
                        m = self._RE_GRAPH_START.match(line)
                        if m:
                            current_pos = 'graph'
                            self.graphs.append(GraceGraph(int(m.group(1))))
                            self.graphs[-1].append(line)
                        else:
                            if self._RE_OBJECT_START.match(line):
                                self.drawing_objects.append(GraceDrawing())
                            self.drawing_objects[-1].append(line)

                elif current_pos == 'region':
                    # regions are followed by regions or graphs
                    if self._RE_GRAPH_START.match(line):
                        current_pos = self._make_graph(line)
                        self.graphs[-1].append(line)
                    else:
                        if self._RE_REGION_START.match(line):
                            self.regions.append(GraceRegion())
                        self.regions[-1].append(line)

                elif current_pos == 'graph':
                    m = self._RE_SET_START.match(line)
                    if m:
                        current_pos = 'set'
                        self.sets.append(GraceSet(int(m.group(1)), int(m.group(2))))
                        self.sets[-1].append(line)
                    else:
                        if self._RE_GRAPH_START.match(line):
                            current_pos = self._make_graph(line)
                        self.graphs[-1].append(line)

                elif current_pos == 'set':
                    m = self._RE_SET_END.match(line)
                    if m:
                        current_pos = 'set_end'
                    self.sets[-1].append(line)

                elif current_pos == 'set_end':
                    m = self._RE_SET_START.match(line)
                    if m:
                        current_pos = 'set'
                        self.sets.append(GraceSet(int(m.group(1)), int(m.group(2))))
                        self.sets[-1].append(line)
                    else:
                        if self._RE_GRAPH_START.match(line):
                            current_pos = self._make_graph(line)
                        else:
                            if GraceGraph._RE_SET_START.match(line):
                                current_pos = 'graph'
                            else:
                                raise ValueError(f'Cannot parse line {line}')

                        self.graphs[-1].append(line)

        print(self.drawing_objects)

    def _make_graph(self, line: str):
        m = self._RE_GRAPH_START.match(line)
        g_idx = int(m.group(1))
        if g_idx < len(self.graphs):
            # this assumes that graphs are ordered in agr file even if missing, e.g. we read gß, g1, g3, ...
            while g_idx != len(self.graphs):
                self.graphs.append(GraceGraph(len(self.graphs)))

        self.graphs.append(GraceGraph(g_idx))

        return 'graph'

    @property
    def colors(self):
        _colors = {}
        for line in self.header:
            m = self._RE_COLOR.match(line)
            if m:
                _colors[int(m['id'])] = (m['disp'], (int(m['red']), int(m['green']), int(m['blue'])))

        return _colors

    def get_color(self, color_num):
        color_num = str(color_num)
        for line in self.header:
            m = self._RE_COLOR.match(line)
            if m:
                if m['id'] == color_num:
                    return m['disp'], (int(m['red']), int(m['green']), int(m['blue']))

        return

    def set_color(self, name: str, rgb: tuple, idx: int = None):
        pos = 'before_colors'
        cnt = 0
        for i, line in enumerate(self.header):
            m = self._RE_COLOR.match(line)
            if m:
                pos = 'in_colors'
                cnt += 1
                if (int(m['red']), int(m['green']), int(m['blue'])) == rgb:
                    logger.info(f'color already defined as {m["disp"]}')
                    return m['id']
                elif m['id'] == idx:
                    self.header[i] = f'@map color {idx} to {rgb}, "{name}"\n'

            elif pos == 'in_colors':
                # landing here for first mismatch after all colors
                if idx is None:
                    if cnt >= 64:
                        logger.warn('Maximum numbers of color reached, color is not used by xmgrace')
                    self.header.insert(i, f'@map color {cnt} to {rgb}, "{name}"\n')
                    return cnt
                else:
                    self.header.insert(i, f'@map color {idx} to {rgb}, "{name}"\n')
                    return idx

    @property
    def fonts(self):
        _fonts = {}
        for line in self.header:
            m = self._RE_FONT.match(line)
            if m:
                _fonts[m['id']] = m['disp']

        return _fonts

    @property
    def size(self):
        for line in self.header:
            m = self._RE_PAGE_SIZE.match(line)
            if m:
                return int(m.group(1)) / 72 * 2.54, int(m.group(2)) / 72 * 2.54

    @size.setter
    def size(self, value: tuple):
        width = value[0] * 72 / 2.54
        height = value[1] * 72 / 2.54

        for i, line in enumerate(self.header):
            m = self._RE_PAGE_SIZE.match(line)
            if m:
                self.header[i] = f'@page size {width:.0f}, {height:.0f}\n'
                break

    def convert(self, value, direction='rel'):
        page_size = min(self.size)

        if direction.startswith('rel'):
            return value / page_size
        elif direction.startswith('abs'):
            return value * page_size
        else:
            raise ValueError(f'Use `abs`, `absolute`, `rel`, or `relative` to convert, found {direction} instead.')

    def get_property(self, g, *args):
        if len(args) == 1:
            return self.graphs[g].get_property(args[0])
        elif len(args) == 2:
            return self.graphs[g].set[args[0]].get_property(args[1])
        else:
            raise TypeError(f'get_property takes two or three arguments ({len(args)+1} given)')

    def set_property(self, g, *args):
        if len(args) == 2:
            self.graphs[g][args[0]] = args[1]
        elif len(args) == 3:
            self.graphs[g].set[args[0]][args[1]] = args[2]
        else:
            raise TypeError(f'get_property takes three or four arguments ({len(args)+1} given)')

    def dataset(self, g, s):
        cmp_to = (int(g), int(s))
        for dataset in self.sets:
            if dataset.graph_set == cmp_to:
                return dataset
        return

    def write(self, fname):
        outfile = pathlib.Path(fname)
        with outfile.open('w') as f:
            f.write(str(self.header))
            f.flush()

            for element in [self.regions, self.drawing_objects, self.graphs, self.sets]:
                for obj in element:
                    f.write(str(obj))
                f.flush()


class GraceDrawing(list):
    type_map = {'text': 'string', 'rectangle': 'box'}

    def __init__(self, draw_type: str):
        super().__init__()


        if draw_type not in ['rectangle', 'line', 'ellipse', 'multipts', 'text']:
            raise ValueError(f'Unknown drawing object {draw_type}')

        draw_type = GraceDrawing.type_map.get(draw_type, draw_type)
        self._type = draw_type

    def _create_box(self):
        text = """\
@with box
@    box on
@    box loctype world
@    box 0.0, 0.0, 1.0, 1.0
@    box linestyle 1
@    box linewidth 1.0
@    box color 1
@    box fill color 1
@    box fill pattern 0'
@box def
"""
        for line in text.split('\n'):
            self.append(line + '\n')

    def _create_ellipse(self):
        text = """\
@with ellipse
@    ellipse on
@    ellipse loctype world
@    ellipse 0.0, 0.0, 1.0, 1.0
@    ellipse linestyle 1
@    ellipse linewidth 1.0
@    ellipse color 1
@    ellipse fill color 1
@    ellipse fill pattern 0
@ellipse def
"""
        for line in text.split('\n'):
            self.append(line + '\n')

    def _create_line(self):
        text = """\
@with string
@    string on
@    string loctype world
@    string 0.0, 0.0
@    string color 1
@    string rot 0
@    string font 0
@    string just 0
@    string char size 1.000000
@    string def ""
"""


    def _create_line(self):
        text = """\
@with line
@    line on
@    line loctype world
@    line g0
@    line 0.0, 0.0, 1.0, 1.0
@    line linewidth 1.0
@    line linestyle 1
@    line color 1
@    line arrow 0
@    line arrow type 0
@    line arrow length 1.000000
@    line arrow layout 1.000000, 1.000000
@line def
"""
        for line in text.split('\n'):
            self.append(line + '\n')



    def __str__(self):
        return ''.join(self)


class GraceHeader(list):
    def __str__(self):
        return ''.join(self)


class GraceProperties(list):
    _RE_ENTRY = re.compile(r'(?!.*(on|off)$)'  # ignore lines that end with on or off
                           r'@\s*(?P<graph>[gs]\d+)?\s*'  # @ maybe followed by g0 or s12 
                           r'(?P<key>[\w\s]*)\s+'   # key: stops at last space unless comma-separated values 
                           r'(?P<val>(?:\s*[\d\w.+-]+\s*,)*\s*[\\{}\"\w.+\- ]+)',  # one value, maybe more with commas
                           re.IGNORECASE | re.VERBOSE)
    _RE_ONOFF = re.compile(r'@\s(?P<graph>[gs]\d+)*\s*(?P<val>[\w\s]*)\s+(on|off)')

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def get_property(self, key: str):
        props = []
        for line in self:
            m = self._RE_ENTRY.match(line)
            if m:
                k = m.group('key').strip()
                if k == key:
                    props.append(_convert_to_value(m.group('val')))

        if len(props) == 1:
            return props[0]
        else:
            return

    def set_property(self, **kwargs):
        raise NotImplementedError

    def _set_property(self, **kwargs):

        for i, line in enumerate(self):
            m = self._RE_ENTRY.match(line)
            if m:
                k = m.group('key')
                if k in kwargs:
                    # re.sub is probably better, but I could not work it out for labels with backslashes (\S)
                    self[i] = line[:m.start('val')] + str(kwargs[k]) + line[m.end('val'):]
                    kwargs.pop(k)

                    if not kwargs:
                        break

        return kwargs

    def get_all_properties(self):
        props = {}
        for line in self:
            m = self._RE_ENTRY.match(line)
            if m:
                key = m.group('key')
                val = _convert_to_value(m.group('val'))
                if key not in props:
                    props[key] = val
                elif isinstance(props[key], tuple):
                    props[key] = props[key] + (val,)
                else:
                    props[key] = props[key], val

        props.pop('with')

        return props

    def get_onoff(self, key: str) -> Optional[bool]:
        onoff_pattern = re.compile(rf'@\s*{key}\s+(?P<val>(on|off))', re.IGNORECASE)
        for i, line in enumerate(self):
            m = onoff_pattern.match(line)

            if m:
                return m.group('val') == 'on'
            return

    def set_onoff(self, key: str, onoff: bool):
        raise NotImplementedError

    def _set_onoff(self, key , onoff: bool):

        bool_str = 'on' if onoff else 'off'
        for i, line in enumerate(self):
            m = key.match(line)
            if m:
                self[i] = line[:m.start('val')] + bool_str + '\n'

                return True

        return False


class GraceGraph(GraceProperties):
    _RE_SET_START = re.compile(r'@\s*s(\d+)\s+\w*', re.IGNORECASE)

    def __init__(self, idx):
        super().__init__()
        self.set = []
        self.idx = int(idx)
        self.__current_pos = 'graph'
        self.__current_idx = None

    def __str__(self):
        return ''.join(self) + ''.join(map(str, self.set))

    def copy(self):
        _cp = GraceGraph(self.idx)
        for line in self:
            _cp.append(line)

        return _cp

    def append(self, line: str):
        m = self._RE_SET_START.match(line)
        if m:
            if self.__current_pos == 'graph' or self.__current_idx != m.group(1):
                self.__current_pos = 'set'
                self.__current_idx = m.group(1)

                while int(self.__current_idx) != len(self.set):
                    self.set.append(GraceSetProps(len(self.set)))

                self.set.append(GraceSetProps(self.__current_idx))

        if self.__current_pos == 'graph':
            super().append(line)
        elif self.__current_pos == 'set':
            self.set[-1].append(line)
        else:
            raise TypeError('Unknown state')

    def new_set(self):
        max_set_idx = -1
        for s in self.set:
            if s is None:
                continue
            max_set_idx = max(max_set_idx, s.idx)
        self.set.append(GraceSetProps(max_set_idx+1))

        return self.set[-1]

    def renumber(self, idx: int):
        number_re = re.compile('g' + str(self.idx), re.IGNORECASE)
        for i, line in enumerate(self):
            m = list(number_re.finditer(line))
            if m:
                self[i] = line[:m[0].start()+1] + str(idx) + line[m[0].end():]

        self.idx = idx

    def set_limits(self, x=None, y=None):
        for i, line in enumerate(self):
            m = self._RE_ENTRY.match(line)
            if m and m.group('key') == 'world':
                _prev_lims = _convert_to_value(m.group('val'))
                if x is None:
                    x = _prev_lims[0]
                if y is None:
                    y = _prev_lims[1]

                self[i] = f'@    world {x[0]}, {y[0]}, {x[1]}, {y[1]}\n'

    def set_log(self, x=None, y=None):
        kwargs = {}
        for ax, label in [(x, 'x'), (y, 'y')]:
            if ax is not None:
                if ax:
                    kwargs[label + 'axes scale'] = 'Logarithmic'
                    kwargs[label + 'axis  ticklabel format'] = 'power'
                    kwargs[label + 'axis  ticklabel prec'] = '0'
                else:
                    kwargs[label + 'axes scale'] = 'Normal'
                    kwargs[label + 'axis  ticklabel format'] = 'general'
                    kwargs[label + 'axis  ticklabel prec'] = '5'

        for i, line in enumerate(self):
            m = self._RE_ENTRY.match(line)
            if m and m.group('key') in kwargs:
                self[i] = re.sub(m.group('val'), kwargs[m.group('key')], line)

    def set_property(self, **kwargs):
        remaining_kw = self._set_property(**kwargs)

        for k, v in remaining_kw.items():
            self.append(f'@ {k} {v}\n')

    def set_axis_property(self, axis: str, **kwargs):
        ax_kwargs = {}
        for k, v in kwargs.items():
            if k in ['invert', 'scale']:
                ax_kwargs[axis + 'axes  ' + k] = v
            else:
                ax_kwargs[axis + 'axis  ' + k] = v

        self.set_property(**ax_kwargs)

    def get_axis_property(self, axis: str, key: str):
        if key in ['invert', 'scale']:
            ax_key = axis + 'axes  ' + key
        else:
            ax_key = axis + 'axis  ' + key

        return self.get_property(ax_key)

    def set_axis_onoff(self, axis: str, key: str, onoff: bool):
        if key in ['invert', 'scale']:
            ax_key = axis + 'axes  ' + key
        else:
            ax_key = axis + 'axis  ' + key

        return self.set_onoff(ax_key, onoff)

    def set_onoff(self, key: str, onoff: bool):
        onoff_pattern = re.compile(rf'@\s*{key}\s+(?P<val>(on|off))', re.IGNORECASE)
        if not self._set_onoff(onoff_pattern, onoff):
            self.append(f'@ {key} {"on" if onoff else "off"}\n')

    def view_to_world(self, pos) -> Tuple[float, float]:
        view = self.get_property('view')
        world = self.get_property('world')

        log_x = self.get_property('xaxes scale').lower()
        if log_x == 'logarithmic':
            x_world = 10**(log10(world[2]/world[0]) / (view[2]-view[0]) * (view[2]-pos[0]) + log10(world[2]))
        else:
            x_world = (world[2]-world[0]) / (view[2]-view[0]) * (view[2]-pos[0]) + world[2]

        log_y = self.get_property('yaxes scale').lower()
        if log_y == 'logarithmic':
            y_world = 10**(log10(world[3]/world[1]) / (view[3]-view[1]) * (view[3]-pos[1]) + log10(world[3]))
        else:
            y_world = (world[3] - world[1]) / (view[3]-view[1]) * (view[3]-pos[1]) + world[3]

        return x_world, y_world
    
    def world_to_view(self, pos) -> Tuple[float, float]:
        view = self.get_property('view')
        world = self.get_property('world')

        view_pos = pos[:]

        log_axes = [self.get_property('xaxes scale'), self.get_property('yaxes scale')]
        for i in [0, 1]:
            if log_axes[i].lower() == 'logarithmic':
                view_pos[i] = (view[i+2]-view[i]) / log10(world[i+2]/world[i]) * log10(pos[i]/world[i]) + view[i]
            else:
                view_pos[i] = (view[i+2]-view[i]) / (world[i+2]-world[i]) * (pos[i]-world[i]) + view[i]

        return view_pos


class GraceSetProps(GraceProperties):

    def __init__(self, idx):
        super().__init__()
        self.idx = int(idx)

    def __str__(self):
        return ''.join(self)

    def set_property(self, **kwargs):
        remaining_kw = self._set_property(**kwargs)

        for k, v in remaining_kw.items():
            self.append(f'@    s{self.idx} {k} {v}\n')

    def set_onoff(self, key: str, onoff: bool):
        onoff_pattern = re.compile(rf'@\s*s\d{1,2}\s+{key}\s+(?P<val>(on|off))', re.IGNORECASE)
        if not self._set_onoff(onoff_pattern, onoff):
            self.append(f'@    s{self.idx} {key} {"on" if onoff else "off"}\n')

    def set_line(self, **kwargs):
        _kwargs = {'line '+k: v for k, v in kwargs.items()}

        self.set_property(**_kwargs)

    def set_symbol(self, **kwargs):
        _kwargs = {'symbol '+k: v for k, v in kwargs.items()}
        if 'symbol' in kwargs:
            _kwargs['symbol'] = kwargs['symbol']
            _kwargs.pop('symbol symbol')

        self.set_property(**_kwargs)


class GraceSet(list):
    _RE_TYPE = re.compile(r'@type (\w+)')
    column_types = {'xy': 2, 'xydx': 3, 'xydy': 3, 'xydxdx': 4, 'xydydy': 4, 'xydxdy': 4, 'xydxdxdydy': 6,
                    'bar': 2, 'bardy': 3, 'bardydy': 4,
                    'xyhilo': 5,  'xyz': 3,
                    'xyr': 3, 'xysize': 3, 'xycolor': 3, 'xycolpat': 4,
                    'xyvmap': 4, 'xyboxplot': 6}

    def __init__(self, g, s):
        super().__init__()
        self._graph = int(g)
        self._set = int(s)

    def __str__(self):
        return ''.join(self)

    @property
    def graph(self):
        return self._graph

    @graph.setter
    def graph(self, new_graph: int):
        self._graph = new_graph
        self[0] = f'@target G{new_graph}.S{self._set}\n'

    @property
    def set(self):
        return self._set

    @set.setter
    def set(self, new_set: int):
        self._set = int(new_set)
        self[0] = f'@target G{self._graph}.S{new_set}\n'

    @property
    def graph_set(self):
        return self._graph, self._set

    @graph_set.setter
    def graph_set(self, new_idx: tuple):
        self._graph, self._set = new_idx
        self[0] = f'@target G{new_idx[0]}.S{new_idx[1]}\n'

    @property
    def type(self):
        m = GraceSet._RE_TYPE.match(self[1])
        if m:
            return m.group(1)

    @type.setter
    def type(self, new_type):
        if new_type not in GraceSet.column_types:
            raise ValueError(f'Unknown plot type {new_type}.')

        self[1] = f'@type {new_type}\n'

    @property
    def rows(self):
        return len(self) - 3

    @property
    def columns(self):
        if self[2] != '&\n':
            return len(self[2].split())
        else:
            return 0

    @property
    def shape(self):
        return self.rows, self.columns

    @property
    def data(self):
        return np.genfromtxt(StringIO(str(self)), skip_header=2, skip_footer=1)

    @data.setter
    def data(self, value):
        if self.columns != 0 and value.shape[1] != self.columns:
            raise TypeError(f'Number of columns is not {self.columns} (given {value.shape[1]})')

        dtype = self.type
        self.clear()
        self.append(f'@target G{self._graph}.S{self._set}\n')
        self.append(f'@type {dtype}\n')
        for i, line in enumerate(value):
            self.append(' '.join(line.astype(str)) + '\n')
        self.append('&\n')


class GraceRegion(list):
    """
    Just another name for a list
    """
    def __str__(self):
        return ''.join(self)


def _convert_to_value(parse_string):
    tuples = parse_string.split(',')
    for i, v in enumerate(tuples):
        v = v.strip()

        if re.match(r'[+-]?\d+(?:.\d+)?(?:[Ee][+-]?\d+)?', v):
            try:
                tuples[i] = int(v)
            except ValueError:
                tuples[i] = float(v)
        elif re.match(r'\".*\"', v):
            tuples[i] = v[1:-1]
        else:
            tuples[i] = v

    if len(tuples) == 1:
        return tuples[0]

    return tuples


def _convert_to_str(value):
    if isinstance(value, (tuple, list)):
        return ', '.join(map(str, value))
    else:
        return str(value)
