
import pandas as pd
import numpy as np
import logging
logging.basicConfig(level=logging.DEBUG, format=' %(asctime)s - %(levelname)s - %(message)s')


def read_csv_chunks_simple(file, passed_df=None, join_how='inner', sep=',', chunksize=10000, dtype=None, index_col=None):
    temp_df = pd.DataFrame()
    if passed_df is not None and join_how == 'inner':
        if join_how == 'inner':
            for c in pd.read_csv(file, sep=sep, chunksize=chunksize, index_col=index_col, dtype=dtype):
                temp_df = temp_df.append(c[c.index.isin(passed_df.index)])
            return passed_df.join(temp_df, how=join_how)
        else:
            for c in pd.read_csv(file, sep=sep, chunksize=chunksize, index_col=index_col, dtype=dtype):
                temp_df = temp_df.append(c)
            return passed_df.join(temp_df, how=join_how)
    else:
        if join_how == 'inner':
            for c in pd.read_csv(file, sep=sep, chunksize=chunksize, index_col=index_col, dtype=dtype):
                temp_df = temp_df.append(c[c.index.isin(passed_df.index)])
            return temp_df
        else:
            for c in pd.read_csv(file, sep=sep, chunksize=chunksize, index_col=index_col, dtype=dtype):
                temp_df = temp_df.append(c)
            return temp_df


# Title Ratings
def clean_ratings(thresh=1000, passed_df=None, join_how='inner'):
    file = '../input/title_ratings.tsv'
    logging.debug('Cleaning {}'.format(file))
    ratings = pd.read_csv(file, delimiter='\t', index_col='tconst')
    ratings = ratings[ratings['numVotes'] >= thresh]
    if passed_df is not None:
        return passed_df.join(ratings, how=join_how)
    else:
        return ratings


# Title Basics
def clean_basics(passed_df=None, join_how='inner'):
    file = '../input/title_basics.tsv'
    logging.debug('Cleaning {}'.format(file))
    data_types = {'titleType': str, 'primaryTitle': str, 'originalTitle': str, 'isAdult': str, 'startYear': str,
                  'endYear': str, 'runtimeMinutes': str, 'genres': str}
    title_basics = pd.DataFrame(columns=data_types.keys())
    if passed_df is not None:
        for c in pd.read_csv(file, delimiter='\t', dtype=data_types, chunksize=10000,
                             index_col='tconst'):
            title_basics = title_basics.append(c[(c['titleType'] == 'movie') & c.index.isin(passed_df.index)])
        return passed_df.join(title_basics, how=join_how)
    else:
        for c in pd.read_csv(file, delimiter='\t', dtype=data_types, chunksize=10000,
                             index_col='tconst'):
            title_basics = title_basics.append(c[(c['titleType'] == 'movie')])
        return title_basics


# Title AKAS
def clean_akas(passed_df=None, join_how='inner'):
    file = '../input/title_akas.tsv'
    logging.debug('Cleaning {}'.format(file))
    data_types = {'ordering': str, 'title': str,	'region': str,	'language': str,
                  'types': str, 'attributes': str, 'isOriginalTitle': str}
    title_akas = pd.DataFrame()
    if passed_df is not None:
        for c in pd.read_csv(file, sep='\t', chunksize=10000, index_col='titleId', dtype=data_types):
            title_akas = title_akas.append(c[(c.index.isin(passed_df.index)) & (c['isOriginalTitle'] == '1')])
        return passed_df.join(title_akas, how=join_how)
    else:
        for c in pd.read_csv(file, sep='\t', chunksize=10000, index_col='titleId', dtype=data_types):
            title_akas = title_akas.append(c[c['isOriginalTitle'] == '1'])
        return title_akas


# Title Principals
def clean_principals(passed_df=None, join_how='inner'):
    file = '../input/title_principals.tsv'
    logging.debug('Cleaning {}'.format(file))
    data_types = {'ordering': str, 'nconst': str,'category': str, 'job': str, 'characters': str}
    title_principals = pd.DataFrame()
    for c in pd.read_csv(file, delimiter='\t', dtype=data_types, chunksize=10000, index_col='tconst'):  # ~35million rows
        if passed_df is not None and join_how == 'inner':
            c = c[c.index.isin(passed_df.index)]
        if c.shape[0] == 0:
            continue
        # 48355 and 48362 are the same // tt0007702
        # Problem: some cast members will end up in different chunk
        # Need to break ordering of cast into a list
        c_out = pd.DataFrame(columns=['tconst', 'principals'])
        t_count = 0
        for t in c.groupby('tconst'):
            try:
                t_temp = [t[0], t[1].set_index('nconst').to_dict(orient='index')]
            except ValueError:  # Actor listed in movie twice
                n_count = 0
                n_out = pd.DataFrame(columns=['nconst', 'orderings'])
                for n in t[1].groupby('nconst'):  # Names are non unique within a movie title
                    n_temp = [n[0], n[1].set_index('ordering').to_dict(orient='index')]
                    n_out.loc[n_count] = n_temp
                    n_count += 1
                t_temp = [t[0], n_out.set_index('nconst').to_dict(orient='index')]
            c_out.loc[t_count] = t_temp
            t_count += 1

        title_principals = title_principals.append(c_out.set_index('tconst'))
    if passed_df is not None:
        return passed_df.join(title_principals, how=join_how)
    else:
        return title_principals


    # Title crew
def clean_crew(passed_df=None, join_how='inner'):
    file = '../input/title_crew.tsv'
    logging.debug('Cleaning {}'.format(file))
    title_crew = pd.DataFrame()
    data_types = {'directors': str, 'writers': str}
    return read_csv_chunks_simple(file, passed_df=passed_df, join_how=join_how, sep='\t',
                           chunksize=10000, dtype=data_types, index_col='tconst')


def write_df(df):
    logging.debug('Writing files...')
    df.to_csv('../output/imdbMerged.tsv', sep='\t')
    pd.Series(df.index).to_csv('../output/imdbMovieCodes.tsv', sep='\t')


if __name__ == "__main__":
    df = clean_ratings(thresh=1000, passed_df=None, join_how='inner')
    df = clean_basics(passed_df=df, join_how='inner')
    df = clean_akas(passed_df=df, join_how='inner')
    df = clean_principals(passed_df=df, join_how='inner')
    df = clean_crew(passed_df=df, join_how='inner')
    write_df(df)

import pandas as pd
import numpy as np
import logging
logging.basicConfig(level=logging.DEBUG, format=' %(asctime)s - %(levelname)s - %(message)s')


def read_csv_chunks_simple(file, passed_df=None, join_how='inner', sep=',', chunksize=10000, dtype=None, index_col=None):
    temp_df = pd.DataFrame()
    if passed_df is not None and join_how == 'inner':
        if join_how == 'inner':
            for c in pd.read_csv(file, sep=sep, chunksize=chunksize, index_col=index_col, dtype=dtype):
                temp_df = temp_df.append(c[c.index.isin(passed_df.index)])
            return passed_df.join(temp_df, how=join_how)
        else:
            for c in pd.read_csv(file, sep=sep, chunksize=chunksize, index_col=index_col, dtype=dtype):
                temp_df = temp_df.append(c)
            return passed_df.join(temp_df, how=join_how)
    else:
        if join_how == 'inner':
            for c in pd.read_csv(file, sep=sep, chunksize=chunksize, index_col=index_col, dtype=dtype):
                temp_df = temp_df.append(c[c.index.isin(passed_df.index)])
            return temp_df
        else:
            for c in pd.read_csv(file, sep=sep, chunksize=chunksize, index_col=index_col, dtype=dtype):
                temp_df = temp_df.append(c)
            return temp_df


# Title Ratings
def clean_ratings(thresh=1000, passed_df=None, join_how='inner'):
    file = '../input/title_ratings.tsv'
    logging.debug('Cleaning {}'.format(file))
    ratings = pd.read_csv(file, delimiter='\t', index_col='tconst')
    ratings = ratings[ratings['numVotes'] >= thresh]
    if passed_df is not None:
        return passed_df.join(ratings, how=join_how)
    else:
        return ratings


# Title Basics
def clean_basics(passed_df=None, join_how='inner'):
    file = '../input/title_basics.tsv'
    logging.debug('Cleaning {}'.format(file))
    data_types = {'titleType': str, 'primaryTitle': str, 'originalTitle': str, 'isAdult': str, 'startYear': str,
                  'endYear': str, 'runtimeMinutes': str, 'genres': str}
    title_basics = pd.DataFrame(columns=data_types.keys())
    if passed_df is not None:
        for c in pd.read_csv(file, delimiter='\t', dtype=data_types, chunksize=10000,
                             index_col='tconst'):
            title_basics = title_basics.append(c[(c['titleType'] == 'movie') & c.index.isin(passed_df.index)])
        return passed_df.join(title_basics, how=join_how)
    else:
        for c in pd.read_csv(file, delimiter='\t', dtype=data_types, chunksize=10000,
                             index_col='tconst'):
            title_basics = title_basics.append(c[(c['titleType'] == 'movie')])
        return title_basics


# Title AKAS
def clean_akas(passed_df=None, join_how='inner'):
    file = '../input/title_akas.tsv'
    logging.debug('Cleaning {}'.format(file))
    data_types = {'ordering': str, 'title': str,	'region': str,	'language': str,
                  'types': str, 'attributes': str, 'isOriginalTitle': str}
    title_akas = pd.DataFrame()
    if passed_df is not None:
        for c in pd.read_csv(file, sep='\t', chunksize=10000, index_col='titleId', dtype=data_types):
            title_akas = title_akas.append(c[(c.index.isin(passed_df.index)) & (c['isOriginalTitle'] == '1')])
        return passed_df.join(title_akas, how=join_how)
    else:
        for c in pd.read_csv(file, sep='\t', chunksize=10000, index_col='titleId', dtype=data_types):
            title_akas = title_akas.append(c[c['isOriginalTitle'] == '1'])
        return title_akas


# Title Principals
def clean_principals(passed_df=None, join_how='inner'):
    file = '../input/title_principals.tsv'
    logging.debug('Cleaning {}'.format(file))
    data_types = {'ordering': str, 'nconst': str,'category': str, 'job': str, 'characters': str}
    title_principals = pd.DataFrame()
    for c in pd.read_csv(file, delimiter='\t', dtype=data_types, chunksize=10000, index_col='tconst'):  # ~35million rows
        if passed_df is not None and join_how == 'inner':
            c = c[c.index.isin(passed_df.index)]
        if c.shape[0] == 0:
            continue
        # 48355 and 48362 are the same // tt0007702
        # Problem: some cast members will end up in different chunk
        # Need to break ordering of cast into a list
        c_out = pd.DataFrame(columns=['tconst', 'principals'])
        t_count = 0
        for t in c.groupby('tconst'):
            try:
                t_temp = [t[0], t[1].set_index('nconst').to_dict(orient='index')]
            except ValueError:  # Actor listed in movie twice
                n_count = 0
                n_out = pd.DataFrame(columns=['nconst', 'orderings'])
                for n in t[1].groupby('nconst'):  # Names are non unique within a movie title
                    n_temp = [n[0], n[1].set_index('ordering').to_dict(orient='index')]
                    n_out.loc[n_count] = n_temp
                    n_count += 1
                t_temp = [t[0], n_out.set_index('nconst').to_dict(orient='index')]
            c_out.loc[t_count] = t_temp
            t_count += 1

        title_principals = title_principals.append(c_out.set_index('tconst'))
    if passed_df is not None:
        return passed_df.join(title_principals, how=join_how)
    else:
        return title_principals


    # Title crew
def clean_crew(passed_df=None, join_how='inner'):
    file = '../input/title_crew.tsv'
    logging.debug('Cleaning {}'.format(file))
    title_crew = pd.DataFrame()
    data_types = {'directors': str, 'writers': str}
    return read_csv_chunks_simple(file, passed_df=passed_df, join_how=join_how, sep='\t',
                           chunksize=10000, dtype=data_types, index_col='tconst')


def write_df(df):
    logging.debug('Writing files...')
    df.to_csv('../output/imdbMerged.tsv', sep='\t')
    pd.Series(df.index).to_csv('../output/imdbMovieCodes.tsv', sep='\t')


if __name__ == "__main__":
    df = clean_ratings(thresh=1000, passed_df=None, join_how='inner')
    df = clean_basics(passed_df=df, join_how='inner')
    df = clean_akas(passed_df=df, join_how='inner')
    df = clean_principals(passed_df=df, join_how='inner')
    df = clean_crew(passed_df=df, join_how='inner')
    write_df(df)
# Title Ratings # Title Basics # Title AKAS # Title Principals # ~35million rows # 48355 and 48362 are the same // tt0007702 # Problem: some cast members will end up in different chunk # Need to break ordering of cast into a list # Actor listed in movie twice # Names are non unique within a movie title # Title crew