# -*- coding: UTF-8 -*-
from __future__ import with_statement
from __future__ import absolute_import
from __future__ import print_function
from fitsrpcserver import Server, deser_array
from visir_utils import Future
import xmlrpclib
import pickle
from threading import Lock
try:
    from contextlib import contextmanager
except ImportError:
    pass

try:
    import multiprocessing as mp
except ImportError:
    # for python < 2.6
    print('import of multiprocessing failed, falling back to threading')
    import threading
    import Queue
    threading.Process = threading.Thread
    # the GIL makes use of more threads pointless
    threading.cpu_count = lambda : 1
    threading.Queue = lambda : Queue.Queue()
    mp = threading

def run_server(port, q):
    try:
        s = Server(port=port)
        q.put(True)
        s.run()
    except:
        q.put(False)

try:
    @contextmanager
    def get_fitsloader():
        f = FitsLoader()
        try:
            yield f
        finally:
            f.shutdown()
except:
    pass

class FitsLoader:
    """ Remote parallel fits loader, distributes open files to remote
        processes to aid parallization of pyfits operations
    """
    def __init__(self):
        self.fnmap = dict()
        # to synchronize the open file dict
        self.lock = Lock()
        try:
            ncpu = min(mp.cpu_count(), 16)
        except:
            ncpu = 2
        self.procs = []
        self.servers = []
        self.cur = 0
        i = 0
        while len(self.procs) < ncpu:
            port = np.random.randint(5000,20000) + i
            q = mp.Queue()
            t = (port, mp.Process(target=run_server, args=(port,q)))
            import sys
            # work around DFS12172
            if not hasattr(sys.stdout, 'flush'):
                sys.stdout.flush = lambda : None
            if not hasattr(sys.stderr, 'flush'):
                sys.stderr.flush = lambda : None
            t[1].start()
            if q.get():
                self.procs.append(t)
            i += 1
            if i > ncpu + 100:
                raise Exception("Could not start fits loading servers")

        for p in self.procs:
            have_proxy = False
            while not have_proxy:
                try:
                    # ServerProxy not threadsafe in >= 2.7
                    self.servers.append((xmlrpclib.ServerProxy('http://localhost:%d' % p[0],
                                                               allow_none=True), Lock()))
                    have_proxy = True
                except:
                    have_proxy = False

    def _conv_idx(self, idx):
        """ fits extension, None means all extensions"""
        if idx is None:
            return idx
        try:
            return [int(x) for x in idx]
        except:
            return int(idx)

    def files(self):
        """ returns all open files """
        return sorted(self.fnmap.keys())

    def _get_server(self, fn):
        with self.lock:
            if fn in self.fnmap:
                return self.fnmap[fn]
            server, slock = self.servers[self.cur]
            self.fnmap[fn] = server, slock
            self.cur += 1
            if self.cur >= len(self.servers):
                self.cur = 0
            return server, slock

    def open(self, files):
        """ open a file or a list of files """
        if isinstance(files, basestring):
            files = [files]
        res = []
        for fn in files:
            server, slock = self._get_server(fn)
            slock.acquire()
            f = Future(server.pyfits_open, args=(fn,),
                       data=slock, callback=lambda x: x.data.release())
            res.append(f)
        [x.get() for x in res]

    def close(self, fn):
        """ close file """
        self.lock.acquire()
        if fn in self.fnmap:
            server, slock = self.fnmap[fn]
            del self.fnmap[fn]
            self.lock.release()
            with slock:
                server.pyfits_close(fn)

    def getheader(self, fn, idx):
        """ get header from extension idx from file fn """
        idx = self._conv_idx(idx)
        server, slock = self._get_server(fn)
        with slock:
            return server.pyfits_getheader(fn, idx)

    def getdata(self, fn, idx):
        """ get data from extension idx from file fn """
        idx = self._conv_idx(idx)
        server, slock = self._get_server(fn)
        with slock:
            data = server.pyfits_getdata(fn, idx)
        if isinstance(idx, int):
            return deser_array(data)
        if isinstance(idx, list) or idx == None:
            return [deser_array(x) for x in data]

    def getval(self, fn, keyword, idx):
        """ get header value for keyword from extension idx from file fn """
        idx = self._conv_idx(idx)
        server, slock = self._get_server(fn)
        with slock:
            data = server.pyfits_getval(fn, keyword, idx)
        if isinstance(idx, int):
            return data
        if isinstance(idx, list) or idx == None:
            return [deser_array(x) for x in data]

    def getval_all(self, keyword, idx=None, files=None, skip_first=True):
        """ get header value for keyword from extension idx from files fn """
        if files is None:
            files = sorted(self.fnmap.keys())
        idx = self._conv_idx(idx)
        res = []

        for fn in files:
            server, slock = self._get_server(fn)
            slock.acquire()
            res.append(Future(server.pyfits_getval, args=(fn, keyword, idx, skip_first),
                              data=slock, callback=lambda x: x.data.release()))
        if isinstance(idx, int):
            return [x.get() for x in res]
        if isinstance(idx, list) or idx == None:
            return [deser_array(x.get()) for x in res]

    def custom(self, fn, fun, args=None):
        """ execute custom function on fits servers
            function receives opened fits file and args as arguments
        """
        server, slock = self._get_server(fn)
        pfun = pickle.dumps(fun)
        with slock:
            return server.pyfits_custom(fn, pfun, args)

    def len(self, fn):
        """ get number of extensions of a fits file """
        return self.custom(fn, len, None)

    def shutdown(self):
        """ shutdown the fits servers """
        for s, slock in self.servers:
            with slock:
                s.quit()
        del self.servers[:]
        for p in self.procs:
            p[1].join()
        del self.procs[:]


import unittest
try:
    import pyfits
except ImportError:
    import astropy.io.fits as pyfits
import numpy as np
import tempfile
class TestFitsLoader(unittest.TestCase):
    def setUp(self):
        fits = []
        self.next = []
        self.img = []
        fns = [tempfile.NamedTemporaryFile() for x in range(10)]
        fns.append(tempfile.NamedTemporaryFile(prefix=u'БВГДЕЁЖЗИЙК'))
        self.fns = fns
        for i, fn in enumerate(fns):
            hdu = pyfits.HDUList()
            hdu.append(pyfits.PrimaryHDU())
            hdu[0].header.update("TEST", i)
            hdu[0].header.update("TEST2", "bla")
            self.next.append(np.random.randint(1, 100))
            self.img.append([])
            for i in range(self.next[-1] - 1):
                img = pyfits.ImageHDU(data=np.random.uniform(size=(50,50)))
                img.header.update("SUB", i)
                self.img[-1].append(img.data)
                hdu.append(img)
            hdu.writeto(fn.name)
        for fn in fns:
            fits.append(pyfits.open(fn))

    def tearDown(self):
        if hasattr(self, "fl"):
            self.fl.shutdown()

    def test_all(self):
        # should use 2.7 setUpClass and be split into several tests
        fns = self.fns
        fl = FitsLoader()
        self.fl = fl
        for fn in fns:
            fl.open(fn.name)

        self.assertEqual(fl.files(), sorted([x.name for x in fns]))

        for i in range(len(self.fns)):
            header = fl.getheader(fns[i].name, 0)
            self.assertTrue(isinstance(header, dict))
            self.assertTrue("TEST" in header)
            self.assertTrue("TEST2" in header)
            self.assertEqual(header["TEST"], i)
            self.assertEqual(header["TEST2"], "bla")
            next = fl.custom(fns[i].name, len)
            self.assertEqual(next, self.next[i])
            for j in range(1, next):
                self.assertTrue((fl.getdata(fns[i].name, j) == self.img[i][j-1]).all())

        idx = range(1, fl.custom(fns[2].name, len))
        for i, r in enumerate(fl.getdata(fns[2].name, idx)):
            self.assertTrue((r == self.img[2][i]).all())

        self.assertEqual(fl.getval(fns[6].name, "TEST2", 0), "bla")

        self.assertEqual(fl.len(fns[-1].name), self.next[-1])

        self.assertEqual(fl.getval_all("TEST2", idx=0), ["bla"] * len(fns))
        self.assertEqual([list(x) for x in fl.getval_all("SUB", idx=None, files=[fns[0].name, fns[1].name])],
                         [range(self.next[0] - 1), range(self.next[1]-1)])

        f = []
        for fn in fns:
            f.append(Future(fl.getval, args=(fn.name, "TEST2", 0)))
        for x in f:
            self.assertEqual(x.get(), "bla")

        fl.close(fns[6].name)
        for fn in fns[:5]:
            fl.close(fn.name)
        fl.shutdown()

    def test_context(self):
        with get_fitsloader() as fl:
            fl.open([x.name for x in self.fns])

if __name__ == "__main__":
    unittest.main()
