# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import csv
def create_csvs(infile):
    pos = infile.rfind('/')
    file = infile[:pos+1] + '/csvs'
    
    filmfile = open(file + "films.csv", "w")
    film_writer = csv.writer(filmfile, delimiter='\t')
    
    edgefile = open(file + "edges.csv", "w")
    edge_writer = csv.writer(edgefile, delimiter='\t')

    directorfile = open(file + "directors.csv", "w")
    director_writer = csv.writer(directorfile, delimiter='\t')
    
    grantfile = open(file + "grants.csv", "w")
    grant_writer = csv.writer(grantfile, delimiter='\t')
    
    countryfile = open(file + "countries.csv", "w")
    country_writer = csv.writer(countryfile, delimiter='\t')

    
    film_header = ["fid:ID", "name", "year", ":LABEL"]
    film_writer.writerow(film_header)    
    
    edge_header = [":START_ID", ":END_ID", ":TYPE"]
    edge_writer.writerow(edge_header)

    director_header = ["dirid:ID", "name", ":LABEL"]
    director_writer.writerow(director_header)
    
    country_header = ["cid:ID", "name", ":LABEL"]
    country_writer.writerow(country_header)
    
    grant_header = ["gid:ID", "name", ":LABEL"]
    grant_writer.writerow(grant_header)

    hdict = {'countries' : {}, 
             'films' : {}, 
             'directors' : {}, 
             'grants' : {}
            }
    
    with open(infile) as csvfile:
        reader = csv.reader(csvfile, delimiter='|')
        for row in reader:
            try:
                if row[0] not in hdict['films'].keys():
                    hdict['films'][row[0]] = ['f' + str(len(hdict['films'].keys()) + 1), row[0].strip(), row[2].strip(), "film"]
                id_film =  hdict['films'][row[0]][0]   

                directores = row[1].split("/")
                countries = row[3].split("/")
                grants  = row[4].split("/") if len(row) > 4  else []

                for d in directores:   
                    d = d.strip()
                    if d not in hdict['directors'].keys():
                        hdict['directors'][d] = ['d' + str(len(hdict['directors'].keys()) + 1), d, "director"]
                    id_dir =  hdict['directors'][d][0]
                    edge_writer.writerow([id_dir, id_film, "DIRECTED"])

                for c in countries:
                    c = c.strip()
                    if c == "2011" or c == "2013":
                        print(row)
                    if c not in hdict['countries'].keys():
                        hdict['countries'][c] = ['c' + str(len(hdict['countries'].keys()) + 1), c, "country"]
                    id_count =  hdict['countries'][c][0]
                    edge_writer.writerow([id_film, id_count, "IS_FROM"])
                    

                for g in grants:
                    g = g.strip()
                    if g not in hdict['grants'].keys():
                        hdict['grants'][g] = ['g' + str(len(hdict['grants'].keys()) + 1), g, "grant"]
                    id_grant =  hdict['grants'][g][0]
                    edge_writer.writerow([id_film, id_grant, "FUNDED_BY"])    
            except:
                print(row)

    for k, v in hdict["countries"].items():
        country_writer.writerow(v)      
        
    for k, v in hdict["directors"].items():
        director_writer.writerow(v)
        
    for k, v in hdict["grants"].items():
        grant_writer.writerow(v)    

    for k, v in hdict["films"].items():
        film_writer.writerow(v)  
        
    filmfile.close()
    edgefile.close()
    directorfile.close()
    countryfile.close()

create_csvs('data/HBF_2002_2018.csv')

# !cp data/csvs2/* /usr/share/neo4j/import/

# !/usr/share/neo4j/bin/neo4j-admin import --mode=csv --database=hbf.db --nodes:film /usr/share/neo4j/import/films.csv --nodes:country /usr/share/neo4j/import/country.csv  --nodes:director /usr/share/neo4j/import/director.csv --relationships /usr/share/neo4j/import/edges.csv --id-type=string --delimiter='\t'

a= ' '
a.strip()



StarcoderdataJupyterScriptsDedupFiltered