# -*- coding: utf-8 -*-
# This file is part of the pyMOR project (http://www.pymor.org).
# Copyright Holders: Rene Milk, Stephan Rave, Felix Schindler
# License: BSD 2-Clause License (http://opensource.org/licenses/BSD-2-Clause)

from __future__ import absolute_import, division, print_function

from itertools import izip
from numbers import Number
import atexit
from collections import OrderedDict
import tempfile
import getpass
import os
import shutil
import numpy as np

from pymor.core.defaults import defaults
from pymor.core.pickle import dump, load
from pymor.la.interfaces import VectorArrayInterface
from pymor.la.listvectorarray import ListVectorArray


_registered_paths = set()
def cleanup():
    for path in _registered_paths:
        shutil.rmtree(path)
atexit.register(cleanup)


@defaults('path')
def basedir(path=os.path.join(tempfile.gettempdir(), 'pymor.diskarray.' + getpass.getuser())):
    if not os.path.exists(path):
        os.mkdir(path)
    _registered_paths.add(path)
    return path


class DiskVectorArray(VectorArrayInterface):
    """|VectorArray| implementation via a list of vectors stored in temporary files."""

    _NONE = tuple()

    def __init__(self, vectors, subtype=_NONE):
        if isinstance(vectors, ListVectorArray):
            assert subtype is self._NONE \
                or (type(subtype) is tuple and len(subtype) == 2
                    and subtype[0] == vectors.vector_type and subtype[1] == vectors.subtype)
            subtype = (vectors.vector_type and vectors.subtype)
            vectors = vectors._list
        else:
            vectors = list(vectors)
        if subtype is self._NONE:
            assert len(vectors) > 0
            subtype = (type(vectors[0]), vectors[0].subtype)
        assert all(type(v) == subtype[0] and v.subtype == subtype[1] for v in vectors)
        self.vector_type, self.vector_subtype = subtype
        self.dir = tempfile.mkdtemp(dir=basedir())
        self._len = len(vectors)
        self.cache_size = 1
        self._cache = OrderedDict()
        for i, v in enumerate(vectors):
            self._store(i, v)

    def __del__(self):
        self.destroy()

    def destroy(self):
        try:
            shutil.rmtree(self.dir)
        except OSError:
            pass

    def _store(self, i, v):
        with open(os.path.join(self.dir, str(i)), 'w') as f:
            dump(v, f)
        self._cache[i] = v
        if len(self._cache) > self.cache_size:
            self._cache.popitem(last=False)

    def _load(self, i):
        if i in self._cache:
            return self._cache[i]
        else:
            with open(os.path.join(self.dir, str(i))) as f:
                v = load(f)
            self._cache[i] = v
            if len(self._cache) > self.cache_size:
                self._cache.popitem(last=False)
            return v

    def get(self, ind=None):
        assert self.check_ind(ind)
        ind = list(xrange(self._len)) if ind is None else [ind] if isinstance(ind, Number) else ind
        return ListVectorArray([self._load(i) for i in ind], subtype=self.subtype, copy=False)

    @classmethod
    def make_array(cls, subtype=None, count=0, reserve=0):
        assert count >= 0
        assert reserve >= 0
        va = cls([], subtype)
        v = va.vector_type.make_zeros(va.vector_subtype)
        for i in range(count):
            va._store(i, v)
        va._len = count
        return va

    def __len__(self):
        return self._len

    @property
    def data(self):
        raise NotImplementedError

    @property
    def dim(self):
        return self.vector_type.make_zeros(self.vector_subtype).dim

    @property
    def subtype(self):
        return (self.vector_type, self.vector_subtype)

    def copy(self, ind=None):
        assert self.check_ind(ind)
        ind = list(xrange(self._len)) if ind is None else [ind] if isinstance(ind, Number) else ind
        c = type(self)([], subtype=self.subtype)
        for d, s in enumerate(ind):
            shutil.copy(os.path.join(self.dir, str(s)),
                        os.path.join(c.dir, str(d)))
        c._len = len(ind)
        return c

    def append(self, other, o_ind=None, remove_from_other=False):
        assert other.check_ind(o_ind)
        assert other.space == self.space
        assert other is not self or not remove_from_other
        o_ind = list(xrange(other._len)) if o_ind is None else [o_ind] if isinstance(o_ind, Number) else o_ind

        self._cache.clear()
        if not remove_from_other:
            for d, s in enumerate(o_ind):
                shutil.copy(os.path.join(other.dir, str(s)),
                            os.path.join(self.dir, str(d + self._len)))
            self._len += len(o_ind)
        else:
            other._cache.clear()
            if len(set(o_ind)) < len(o_ind):
                self.append(other, o_ind=o_ind, remove_from_other=False)
                other.remove(o_ind)
            else:
                for d, s in enumerate(o_ind):
                    shutil.move(os.path.join(other.dir, str(s)),
                                os.path.join(self.dir, str(d + self._len)))
                self._len += len(o_ind)

                remaining = sorted(set(xrange(other._len)) - set(o_ind))
                for d, s in enumerate(remaining):
                    shutil.move(os.path.join(other.dir, str(s)),
                                os.path.join(other.dir, str(d)))
                other._len = len(remaining)

    def remove(self, ind=None):
        assert self.check_ind(ind)
        self._cache.clear()
        ind = list(xrange(self._len)) if ind is None else [ind] if isinstance(ind, Number) else list(set(ind))
        for i in ind:
            os.remove(os.path.join(self.dir, str(i)))
        remaining = sorted(set(xrange(len(self))) - set(ind))
        for d, s in enumerate(remaining):
            shutil.move(os.path.join(self.dir, str(s)),
                        os.path.join(self.dir, str(d)))
        self._len = len(remaining)

    def replace(self, other, ind=None, o_ind=None, remove_from_other=False):
        assert self.check_ind_unique(ind)
        self._cache.clear()
        assert other.check_ind(o_ind)
        assert other.space == self.space
        assert other is not self or not remove_from_other
        ind = list(xrange(self._len)) if ind is None else [ind] if isinstance(ind, Number) else ind
        o_ind = list(xrange(other._len)) if o_ind is None else [o_ind] if isinstance(o_ind, Number) else o_ind
        assert len(ind) == len(o_ind)

        if other is self:
            backup = set(ind).intersection(set(o_ind))
            for i in backup:
                shutil.move(os.path.join(self.dir, str(i)),
                            os.path.join(self.dir, '_' + str(i)))
            for d, s in izip(ind, o_ind):
                if s in backup:
                    shutil.copy(os.path.join(self.dir, '_' + str(s)),
                                os.path.join(self.dir, str(d)))
                else:
                    shutil.copy(os.path.join(self.dir, str(s)),
                                os.path.join(self.dir, str(d)))
            for i in backup:
                os.remove(os.path.join(self.dir, '_' + str(i)))
        else:
            for d, s in izip(ind, o_ind):
                shutil.copy(os.path.join(other.dir, str(s)),
                            os.path.join(self.dir, str(d)))
            if remove_from_other:
                other.remove(o_ind)

    def almost_equal(self, other, ind=None, o_ind=None, rtol=None, atol=None):
        assert self.check_ind(ind)
        assert other.check_ind(o_ind)
        assert other.space == self.space
        ind = list(xrange(self._len)) if ind is None else [ind] if isinstance(ind, Number) else ind
        o_ind = list(xrange(other._len)) if o_ind is None else [o_ind] if isinstance(o_ind, Number) else o_ind

        if len(ind) == 1:
            a = self._load(ind[0])
            return np.array([a.almost_equal(other._load(oi), rtol=rtol, atol=atol) for oi in o_ind])
        elif len(o_ind) == 1:
            b = other._load(o_ind[0])
            return np.array([self._load(i).almost_equal(b, rtol=rtol, atol=atol) for i in ind])
        else:
            assert len(ind) == len(o_ind)
            return np.array([self._load(i).almost_equal(other._load(oi), rtol=rtol, atol=atol)
                            for i, oi in izip(ind, o_ind)])

    def scal(self, alpha, ind=None):
        assert self.check_ind_unique(ind)
        assert isinstance(alpha, Number) \
            or isinstance(alpha, np.ndarray) and alpha.shape == (self.len_ind(ind),)
        ind = list(xrange(self._len)) if ind is None else [ind] if isinstance(ind, Number) else ind
        if isinstance(alpha, np.ndarray):
            for a, i in izip(alpha, ind):
                new_vec = self._load(i)
                new_vec.scal(a)
                self._store(i, new_vec)
        else:
            for i in ind:
                new_vec = self._load(i)
                new_vec.scal(alpha)
                self._store(i, new_vec)

    def axpy(self, alpha, x, ind=None, x_ind=None):
        assert self.check_ind_unique(ind)
        assert x.check_ind(x_ind)
        assert self.space == x.space
        assert self.len_ind(ind) == x.len_ind(x_ind) or x.len_ind(x_ind) == 1
        assert isinstance(alpha, Number) \
            or isinstance(alpha, np.ndarray) and alpha.shape == (self.len_ind(ind),)
        ind = list(xrange(self._len)) if ind is None else [ind] if isinstance(ind, Number) else ind
        x_ind = list(xrange(x._len)) if x_ind is None else [x_ind] if isinstance(x_ind, Number) else x_ind

        if self is x:
            if set(ind).intersection(set(x_ind)):
                self.axpy(alpha, x.copy(x_ind), ind)
                return

        if np.all(alpha == 0):
            return

        if len(x_ind) == 1:
            x = x._load(x_ind[0])
            if isinstance(alpha, np.ndarray):
                for a, y in izip(alpha, ind):
                    new_vec = self._load(y)
                    new_vec.axpy(a, x)
                    self._store(y, new_vec)
            else:
                for y in ind:
                    new_vec = self._load(y)
                    new_vec.axpy(alpha, x)
                    self._store(y, new_vec)
        else:
            if isinstance(alpha, np.ndarray):
                for a, xx, y in izip(alpha, x_ind, ind):
                    new_vec = self._load(y)
                    new_vec.axpy(a, x._load(xx))
                    self._store(y, new_vec)
            else:
                for xx, y in izip(x_ind, ind):
                    new_vec = self._load(y)
                    new_vec.axpy(alpha, x._load(xx))
                    self._store(y, new_vec)

    def dot(self, other, pairwise, ind=None, o_ind=None):
        assert self.check_ind(ind)
        assert other.check_ind(o_ind)
        assert self.space == other.space
        ind = list(xrange(self._len)) if ind is None else [ind] if isinstance(ind, Number) else ind
        o_ind = list(xrange(other._len)) if o_ind is None else [o_ind] if isinstance(o_ind, Number) else o_ind