
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm
 import scoped_session, sessionmaker, mapper
from sqlalchemy.engine import Engine
from sqlalchemy import event, create_engine, func

from sqlalchemy import and_, or_, schema, types, MetaData, Sequence, Column, ForeignKey, UniqueConst
raint, Table
from sqlalchemy.orm import relationship, backref, dynamic_loader, deferred, reconstructor
from sqlalchemy.orm.collections import column_mapped_collection, attribute_mapped_collection
from sqlalchemy.orm.interfaces import MapperExtension
from sqlalchemy.orm.exc import NoResultFound
from sqlalchemy.orm.session import object_session
from sqlalchemy.exc import OperationalError, IntegrityError
from sqlalchemy.ext.associationproxy import association_proxy
from sqlalchemy.ext.declarative import declared_attr, declarative_base
from sqlalchemy.sql.functions import current_timestamp

from zope.sqlalchemy import ZopeTransactionExtension

from fatools.lib.utils import cout, cerr
from fatools.lib.fautil.mixin2 import ( PanelMixIn, FSAMixIn, ChannelMixIn, MarkerMixIn,
                BinMixIn, AlleleSetMixIn, AlleleMixIn, SampleMixIn, BatchMixIn,
                NoteMixIn, BatchNoteMixIn, SampleNoteMixIn, FSANoteMixIn,
                ChannelNoteMixIn, AlleleSetNoteMixIn, PanelNoteMixIn, MarkerNoteMixIn )

import os, io, yaml

#__all__ = ['get_base', 'get_dbsession', 'set_datalogger']

# this is necessary for SQLite to use FOREIGN KEY support (as well as ON DELETE CASCADE)
@event.listens_for(Engine, 'connect')
def set_sqlite_pragma(dbapi_connection, connection_record):
    cursor = dbapi_connection.cursor()
    cursor.execute('PRAGMA foreign_keys=ON')
    cursor.close()


Base = declarative_base()

def _generic_query(cls, session):
    return session.query(cls)

Base.query = classmethod(_generic_query)

def _generic_get(cls, pkid, session):
    q = cls.query(session)
    return q.get(int(pkid))

Base.get = classmethod(_generic_get)

def _generic_lowername(cls):
    return cls.__name__.lower()

Base.lowername = classmethod(_generic_lowername)

def _generic_delete(cls, pkid, session):
    q = cls.query(session)
    return q.filter(cls.id == int(pkid)).delete()

Base.delete = classmethod(_generic_delete)


## JSONCol and YAMLCol are taken from Rhombus

# create JSON column
# XXX: may be more appropriate to create a dict-based object that will serialize to JSON?

null = object()

class JSONCol(types.TypeDecorator):
    impl = types.Unicode

    def process_bind_param(self, value, dialect):
        if value is null:
            value = None
        return json.dumps(value)

    def process_result_value(self, value, dialect):
        if value is None:
            return None
        return json.loads(value)

    def copy_value(self, value):
        return copy.deepcopy(value)

# create YAML column
#

null = null

class YAMLCol(types.TypeDecorator):
    impl = types.Unicode

    def process_bind_param(self, value, dialect):
        if value is null:
            value = None
        return yaml.dump(value, default_flow_style=True)

    def process_result_value(self, value, dialect):
        if value is None:
            return None
        return yaml.load(value)

    def copy_value(self, value):
        return copy.deepcopy(value)

import numpy, copy

class NPArray(types.TypeDecorator):
    impl = types.LargeBinary

    def process_bind_param(self, value, dialect):
        if value is None:
            return None
        #buf = value.tostring()
        buf = io.BytesIO()
        numpy.save(buf, value)
        return buf.getvalue()

    def process_result_value(self, value, dialect):
        if value is None:
            return None
        buf = io.BytesIO(value)
        return numpy.load(buf)
        #return numpy.fromstring(value)

    def copy_value(self, value):
        return copy.deepcopy( value )



class Note(Base, NoteMixIn):

    __tablename__ = 'notes'
    id = Column(types.Integer, primary_key=True)
    text = Column(types.String(1024), nullable=False, default='')
    cat = Column(types.String(32), nullable=False, default='')
    stamp = Column(types.DateTime, nullable=False)



class Batch(Base, BatchMixIn):

    __tablename__ = 'batches'
    id = Column(types.Integer, primary_key=True)
    code = Column(types.String(16), nullable=False, unique=True)
    fsa_provider = Column(types.String(32), nullable=False)
    species = Column(types.String(16), nullable=False)
    description = Column(types.String(1024), nullable=False, default='')
    remark = deferred(Column(types.String(1024), nullable=True))
    data = deferred(Column(YAMLCol(4096), nullable=False, default=''))
    bin_batch_id = Column(types.Integer, ForeignKey('batches.id'), nullable=True)

    bin_batch = relationship('Batch', uselist=False)


    def add_sample(self, sample_code):
        """ return a new Sample with sample_code """

        sample = Sample()
        sample.code = sample_code
        sample.batch = self

        return sample


    def search_sample(self, sample_code):
        """ return a single Sample from the current batch with sample_code """
        try:
            return self.samples.filter( func.lower(Sample.code) == func.lower(sample_code),
                Sample.batch_id == self.id ).one()
        except NoResultFound:
            return None


    @staticmethod
    def search(code, session):
        """ provide case-insensitive search for batch code """
        q = Batch.query(session).filter( func.lower(Batch.code) == func.lower(code) )
        return q.one()


    def get_panel(self, panel_code):
        return Panel.search(panel_code, object_session(self))


    def get_marker(self, marker_code, species=None):
        session = object_session(self)
        markers = None  #XXX: Fix me
        raise NotImplementedError()


    @property
    def sample_ids(self):
        """ faster implementation of getting sample ids """
        session = object_session(self)
        return [ x[0] for x in session.query( Sample.id ).filter( Sample.batch_id == self.id ) ]



class BatchNote(Base, BatchNoteMixIn):

    __tablename__ = 'batchnotes'
    id = Column(types.Integer, primary_key=True)
    batch_id = Column(types.Integer, ForeignKey('batches.id', ondelete='CASCADE'),
                nullable=False)
    note_id = Column(types.Integer, ForeignKey('notes.id', ondelete='CASCADE'),
                nullable=False)



class Sample(Base, SampleMixIn):

    __tablename__ = 'samples'
    id = Column(types.Integer, primary_key=True)
    code = Column(types.String(64), nullable=False)
    type = Column(types.String(1), default='S')
    altcode = Column(types.String(16), nullable=True)               # custom usage
    category = Column(types.Integer, nullable=False, default=0)     # custom usage
    batch_id = Column(types.Integer, ForeignKey('batches.id', ondelete='CASCADE'),
                nullable=False)
    int1 = Column(types.Integer, nullable=False, default=-1)        # custom usage
    int2 = Column(types.Integer, nullable=False, default=-1)        # custom usage
    string1 = Column(types.String(16), nullable=False, default='')  # custom usage
    string2 = Column(types.String(16), nullable=False, default='')  # custom usage
    batch = relationship(Batch, uselist=False,
            backref=backref('samples', lazy='dynamic', cascade='save-update,delete'))
    remark = deferred(Column(types.String(1024), nullable=True))

    __table_args__ = (  UniqueConstraint( 'code', 'batch_id' ),
                        UniqueConstraint( 'altcode', 'batch_id')
                    )

    def new_fsa(self, raw_data, filename, status, panel=None):
        fsa = FSA( raw_data = raw_data, filename = filename )
        if panel is None:
            panel = Panel.search('undefined', object_session(self))
        fsa.panel = panel
        fsa.sample = self
        fsa.status = status
        return fsa



class SampleNote(Base, SampleNoteMixIn):

    __tablename__ = 'samplenotes'
    id = Column(types.Integer, primary_key=True)
    sample_id = Column(types.Integer, ForeignKey('samples.id', ondelete='CASCADE'),
                nullable=False)
    note_id = Colu