import csv
import copy
from datetime import datetime
from database import Database
from data.indicators.sma import calculate_sma
from model.sma import Sma
from model.ticker import Ticker


database = Database()


def load_database(chunk_size=100000):
    database.drop_tables()
    database.create_tables()

    with open("./data/output/tickers.csv") as file:
        try:
            rows = []
            reader = csv.DictReader(file)
            total = 0
            for row in reader:
                ticker = Ticker(
                    date=datetime.strptime(row.get("date"), "%Y-%m-%d").date(),
                    symbol=row.get("symbol"),
                    open=row.get("open"),
                    high=row.get("high"),
                    low=row.get("low"),
                    close=row.get("close"),
                    volume=row.get("volume")
                )
                rows.append(ticker)

                if len(rows) >= chunk_size:
                    database.session.bulk_save_objects(rows)
                    database.session.commit()
                    total += chunk_size
                    print(f"{total} rows inserted")
                    rows = []
            
            # Insert any remaining rows.
            if rows:
                database.session.bulk_save_objects(rows)
                database.session.commit()
                total += len(rows)
                print(f"{total} inserted")
        except Exception as e:
            import pprint
            pprint.pprint(e)


def create_sma():
    from sqlalchemy import distinct

    database.drop_table("sma")
    database.create_tables()

    sma_periods = [5, 10, 20, 50, 100, 200]
    symbols = database.session.query(distinct(Ticker.symbol)).all()
    count = 0
    for symbol in symbols:
        try:
            daily_ticker_data = database.session.query(Ticker).filter_by(symbol=symbol[0]).all()
            daily_ticker_data = [ticker.to_dict() for ticker in daily_ticker_data]
            
            smas = {}
            for period in sma_periods:
                data = copy.deepcopy(daily_ticker_data)
                sma_data = calculate_sma(data, period)
                smas[period] = sma_data
            
            for i in range(len(daily_ticker_data)):
                row = Sma(
                    date=daily_ticker_data[i]['date'],
                    symbol=symbol[0],
                    sma_5=smas[5][i]['SMA'],
                    sma_10=smas[10][i]['SMA'],
                    sma_20=smas[20][i]['SMA'],
                    sma_50=smas[50][i]['SMA'],
                    sma_100=smas[100][i]['SMA'],
                    sma_200=smas[200][i]['SMA']
                )
                database.session.add(row)
            database.session.commit()
            count += 1
            print(f"finished {symbol[0]}: {count} of {len(symbols)}: time: {datetime.now()}")
        except Exception as e:
            print(e)
            database.session.rollback()


if __name__ == "__main__":
    create_sma()