import os
import tempfile
import unittest

import numpy as np
import pandas as pd
from sqlalchemy import create_engine

from tests.settings import POSTGRESQL_ENGINE, SQLITE_ENGINE
from tests.utils import get_repository_path, DBTest
from ukbrest.common.pheno2sql import Pheno2SQL


class Pheno2SQLTest(DBTest):
    @unittest.skip('sqlite being removed')
    def test_sqlite_default_values(self):
        # Prepare
        csv_file = get_repository_path('pheno2sql/example01.csv')
        db_engine = SQLITE_ENGINE

        p2sql = Pheno2SQL(csv_file, db_engine)

        # Run
        p2sql.load_data()

        # Validate
        assert p2sql.db_type == 'sqlite'

        ## Check table exists
        tmp = pd.read_sql("SELECT name FROM sqlite_master WHERE type='table' AND name='{}';".format('ukb_pheno_0_00'), create_engine(db_engine))
        assert not tmp.empty

        ## Check columns are correct
        tmp = pd.read_sql('select * from ukb_pheno_0_00', create_engine(db_engine))
        expected_columns = ["eid","c21_0_0","c21_1_0","c21_2_0","c31_0_0","c34_0_0","c46_0_0","c47_0_0","c48_0_0"]
        assert len(tmp.columns) == len(expected_columns)
        assert all(x in expected_columns for x in tmp.columns)

        ## Check data is correct
        tmp = pd.read_sql('select * from ukb_pheno_0_00', create_engine(db_engine), index_col='eid')
        assert not tmp.empty
        assert tmp.shape[0] == 2
        assert tmp.loc[1, 'c21_0_0'] == 'Option number 1'
        assert tmp.loc[1, 'c21_1_0'] == 'No response'
        assert tmp.loc[1, 'c21_2_0'] == 'Yes'
        assert tmp.loc[1, 'c31_0_0'] == '2012-01-05'
        assert int(tmp.loc[1, 'c34_0_0']) == 21
        assert int(tmp.loc[1, 'c46_0_0']) == -9
        assert tmp.loc[1, 'c47_0_0'].round(5) == 45.55412
        assert tmp.loc[1, 'c48_0_0'] == '2011-08-14'
        assert tmp.loc[2, 'c21_0_0'] == 'Option number 2'
        assert pd.isnull(tmp.loc[2, 'c21_1_0'])
        assert tmp.loc[2, 'c21_2_0'] == 'No'
        assert tmp.loc[2, 'c31_0_0'] == '2015-12-30'
        assert int(tmp.loc[2, 'c34_0_0']) == 12
        assert int(tmp.loc[2, 'c46_0_0']) == -2
        assert tmp.loc[2, 'c47_0_0'].round(5) == -0.55461
        assert tmp.loc[2, 'c48_0_0'] == '2010-03-29'

    def test_postgresql_default_values(self):
        # Prepare
        csv_file = get_repository_path('pheno2sql/example01.csv')
        db_engine = POSTGRESQL_ENGINE

        p2sql = Pheno2SQL(csv_file, db_engine)

        # Run
        p2sql.load_data()

        # Validate
        assert p2sql.db_type == 'postgresql'

        ## Check table exists
        table = pd.read_sql("SELECT EXISTS (SELECT 1 FROM pg_tables WHERE schemaname = 'public' AND tablename = '{}');".format('ukb_pheno_0_00'), create_engine(db_engine))
        assert table.iloc[0, 0]

        ## Check columns are correct
        tmp = pd.read_sql('select * from ukb_pheno_0_00', create_engine(db_engine))
        expected_columns = ["eid","c21_0_0","c21_1_0","c21_2_0","c31_0_0","c34_0_0","c46_0_0","c47_0_0","c48_0_0"]
        assert len(tmp.columns) == len(expected_columns)
        assert all(x in expected_columns for x in tmp.columns)

        ## Check data is correct
        tmp = pd.read_sql('select * from ukb_pheno_0_00', create_engine(db_engine), index_col='eid')
        assert not tmp.empty
        assert tmp.shape[0] == 2
        assert tmp.loc[1, 'c21_0_0'] == 'Option number 1'
        assert tmp.loc[1, 'c21_1_0'] == 'No response'
        assert tmp.loc[1, 'c21_2_0'] == 'Yes'
        assert tmp.loc[1, 'c31_0_0'].strftime('%Y-%m-%d') == '2012-01-05'
        assert int(tmp.loc[1, 'c34_0_0']) == 21
        assert int(tmp.loc[1, 'c46_0_0']) == -9
        assert tmp.loc[1, 'c47_0_0'].round(5) == 45.55412
        assert tmp.loc[1, 'c48_0_0'].strftime('%Y-%m-%d') == '2011-08-14'
        assert tmp.loc[2, 'c21_0_0'] == 'Option number 2'
        assert pd.isnull(tmp.loc[2, 'c21_1_0'])
        assert tmp.loc[2, 'c21_2_0'] == 'No'
        assert tmp.loc[2, 'c31_0_0'].strftime('%Y-%m-%d') == '2015-12-30'
        assert int(tmp.loc[2, 'c34_0_0']) == 12
        assert int(tmp.loc[2, 'c46_0_0']) == -2
        assert tmp.loc[2, 'c47_0_0'].round(5) == -0.55461
        assert tmp.loc[2, 'c48_0_0'].strftime('%Y-%m-%d') == '2010-03-29'

    def test_exit(self):
        # Prepare
        csv_file = get_repository_path('pheno2sql/example01.csv')
        db_engine = POSTGRESQL_ENGINE
        temp_dir = tempfile.mkdtemp()

        # Run
        with Pheno2SQL(csv_file, db_engine, tmpdir=temp_dir) as p2sql:
            p2sql.load_data()

        # Validate
        ## Check table exists
        table = pd.read_sql("SELECT EXISTS (SELECT 1 FROM pg_tables WHERE schemaname = 'public' AND tablename = '{}');".format('ukb_pheno_0_00'), create_engine(db_engine))
        assert table.iloc[0, 0]

        ## Check columns are correct
        tmp = pd.read_sql('select * from ukb_pheno_0_00', create_engine(db_engine))
        expected_columns = ["eid","c21_0_0","c21_1_0","c21_2_0","c31_0_0","c34_0_0","c46_0_0","c47_0_0","c48_0_0"]
        assert len(tmp.columns) == len(expected_columns)
        assert all(x in expected_columns for x in tmp.columns)

        ## Check data is correct
        tmp = pd.read_sql('select * from ukb_pheno_0_00', create_engine(db_engine), index_col='eid')
        assert not tmp.empty
        assert tmp.shape[0] == 2

        ## Check that temporary files were deleted
        assert len(os.listdir(temp_dir)) == 0

    @unittest.skip('sqlite being removed')
    def test_sqlite_less_columns_per_table(self):
        # Prepare
        csv_file = get_repository_path('pheno2sql/example01.csv')
        db_engine = SQLITE_ENGINE

        p2sql = Pheno2SQL(csv_file, db_engine, n_columns_per_table=3)

        # Run
        p2sql.load_data()

        # Validate
        assert p2sql.db_type == 'sqlite'

        ## Check tables exist
        table = pd.read_sql("SELECT name FROM sqlite_master WHERE type='table' AND name='{}';".format('ukb_pheno_0_00'), create_engine(db_engine))
        assert not table.empty

        table = pd.read_sql("SELECT name FROM sqlite_master WHERE type='table' AND name='{}';".format('ukb_pheno_0_01'), create_engine(db_engine))
        assert not table.empty

        table = pd.read_sql("SELECT name FROM sqlite_master WHERE type='table' AND name='{}';".format('ukb_pheno_0_02'), create_engine(db_engine))
        assert not table.empty

        ## Check columns are correct
        tmp = pd.read_sql('select * from ukb_pheno_0_00', create_engine(db_engine))
        expected_columns = ["eid","c21_0_0","c21_1_0","c21_2_0"]
        assert len(tmp.columns) == len(expected_columns)
        assert all(x in expected_columns for x in tmp.columns)

        tmp = pd.read_sql('select * from ukb_pheno_0_01', create_engine(db_engine))
        expected_columns = ["eid","c31_0_0","c34_0_0","c46_0_0"]
        assert len(tmp.columns) == len(expected_columns)
        assert all(x in expected_columns for x in tmp.columns)

        tmp = pd.read_sql('select * from ukb_pheno_0_02', create_engine(db_engine))
        expected_columns = ["eid","c47_0_0","c48_0_0"]
        assert len(tmp.columns) == len(expected_columns)
        assert all(x in expected_columns for x in tmp.columns)

        ## Check data is correct
        tmp = pd.read_sql('select * from ukb_pheno_0_00', create_engine(db_engine), index_col='eid')
        assert not tmp.empty
        assert tmp.shape[0] == 2
        assert tmp.loc[1, 'c21_0_0'] == 'Option number
 1'
        assert tmp.loc[1, 'c21_1_0'] == 'No response'
        assert tmp.loc[1, 'c21_2_0'] == 'Yes'
        assert tmp.loc[2, 'c21_0_0']
 == 'Option number 2'
        assert pd.isnull(tmp.loc[2, 'c21_1_0'])
        assert tmp.loc[2, 'c21_2_0'] == 'No'

        tmp = pd.read_sql('select * from ukb_pheno_0_01', create_engine(db_engine), index_col='eid')
        assert not tmp.empty
        assert tmp.shape[0] == 2
        assert tmp.loc[1, 'c31_0_0'] == '2012-01-05'
        assert int(tmp.loc[1, 'c34_0_0']) == 21
        assert int(tmp.loc[1, 'c46_0_0']) == -9
        assert tmp.loc[2, 'c31_0_0'] == '2015-12-30'
        assert int(tmp.loc[2, 'c34_