import argparse
import csv
import gzip
import logging
import os
import pathlib
import sqlite3
import sys
import time

import pandas as pd

from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker

from ..db import initialise_db
from ..db.models import db
from ..db.models.meta import Base


logging.basicConfig(
    level=logging.INFO,
    format='[%(asctime)s] %(message)s'
)


def get_database(db_path):
    engine = create_engine(f"sqlite:///{db_path}")

    db_session = scoped_session(
        sessionmaker(
            autocommit=False,
            autoflush=False,
            enable_baked_queries=True,
            bind=engine
        )
    )

    Base.query = db_session.query_property()

    return engine, db_session


def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--input_dir", "-i", type=str, default=".")
	ap.add_argument("--output_prefix", "-o", type=str, default="collated")
	ap.add_argument("--column", "-c", type=str)
	ap.add_argument("--db_path", type=str, default="observations.sqlite3")
	
	args = ap.parse_args()

	outdir = os.path.dirname(args.output_prefix)
	if outdir and outdir != ".":
		pathlib.Path(outdir).mkdir(exist_ok=True, parents=True)

	dirpath, _, files = next(os.walk(args.input_dir))

	engine, db_session = get_database(args.db_path)
	initialise_db(engine)

	features_d = {}
	category_id = None

	for i, f in enumerate(files, start=1):
		logging.info(f"Processing file {i}/{len(files)}: {f}")

		t0 = time.time()
		fname = os.path.basename(f).replace(".txt.gz", "")
		*sample, category = fname.split(".")
		sample = ".".join(sample)

		if category_id is None:
			db_category = db.Category(id=0, name=category)
			db_session.add(db_category)
			category_id = db_category.id
		
		db_sample = db.Sample(name=sample)
		db_session.add(db_sample)
		
		counts = pd.read_csv(
			os.path.join(dirpath, f),
			delimiter="\t",
			index_col=0,
			usecols=(("feature", "gene")[category == "gene_counts"], args.column),
		)
		
		for _, row in counts.iterrows():
			if row.name != "unannotated":
				feature_name = row.name
				feature_id = features_d.setdefault(feature_name, len(features_d))
				
				db_observation = db.Observation(
					metric=args.column,
					value=float(row[0]),
					category_id=0,
					sample_id=db_sample.id,
					feature_id=feature_id,
				)
				db_session.add(db_observation)

		db_session.commit()
		logging.info(f"Finished in {time.time() - t0}s.")


	logging.info(f"Adding {len(features_d)} features to database...")
	t0 = time.time()
	for feature_name, feature_id in features_d.items():
		# logging.info(f"    {feature_name} -> {feature_id}")
		db_feature = db.Feature(id=feature_id, name=feature_name, category_id=0)
		db_session.add(db_feature)

	db_session.commit()
	logging.info(f"Finished in {time.time() - t0}s.")	

	features_d.clear()

	logging.info("Converting database to count matrix...")
	t0 = time.time()
	con = sqlite3.connect(args.db_path)
	df = pd.read_sql_query(
		"select sample.name as sample_id, "
		"feature.name as feature_id, "
		"observation.value "
		"from observation "
		"join sample on sample.id = observation.sample_id "
		"join feature on feature.id = observation.feature_id",
		con
	)

	df = df.pivot(index="feature_id", columns="sample_id")
	df.columns = [x[1] for x in df.columns]
	df.index.name = "feature"
	logging.info(f"Finished in {time.time() - t0}s.")
	logging.info("Saving database...")
	t0 = time.time()
	df.to_csv(
		f"{args.output_prefix}.{category}.{args.column}.txt.gz",
		sep="\t",
		na_rep="NA",
	)
	logging.info(f"Finished in {time.time() - t0}s.")


	
	
if __name__ == "__main__":
	main()