# Run this cell to set up the notebook.
import numpy as np
import pandas as pd
import seaborn as sns

import matplotlib
%matplotlib inline
import matplotlib.pyplot as plt

from client.api.notebook import Notebook
ok = Notebook('lab05.ok')

"""
Explanation: Lab 5: Relational Algebra in Pandas
End of explanation
"""


young_sailors = pd.DataFrame({
        "sid":      [2701, 18869, 63940, 21869, 17436],
        "sname":    ["Jerry", "Morgan", "Danny", "Jack", "Dustin"],
        "rating":   [8, 6, 4, 9, 3],
        "age":      [25, 26, 21, 27, 22],
    })
salty_sailors = pd.DataFrame({
        "sid":      [2701, 17436, 45433, 22689, 46535],
        "sname":    ["Jerry", "Dustin", "Balon", "Euron", "Victarion"],
        "rating":   [8, 3, 7, 10, 2],
        "age":      [25, 22, 39, 35, 37],
    })
boats = pd.DataFrame({
        "bid":      [41116, 54505, 50041, 35168, 58324],
        "bname":    ["The Black Sparrow", "The Great Kraken", "The Prophetess", "Silence", "Iron Victory"],
        "color":    ["Black", "Orange", "Silver", "Red", "Grey"],
    })
reservations = pd.DataFrame({
        "sid":      [21869, 45433, 18869, 22689, 21869, 17436, 63940, 45433, 21869, 18869],
        "bid":      [41116, 35168, 50041, 41116, 58324, 50041, 54505, 41116, 50041, 41116],
        "day":      ["3/1", "3/1", "3/2", "3/2", "3/2", "3/3", "3/3", "3/3", "3/3", "3/4"],
    })

"""
Explanation: Boat Club
The Berkeley Boat Club wants to better organize their user data, and they've hired you to do it. Your first job is to implement code for relational algebra operators in python (unlike you, they don't know how to use pandas).
You may want to refer to these slides, to remember what each operation does. You may also want to refer to the pandas documentation.
Here are the Boat Club's databases. Your job is to implement a variety of unary and binary relational algebra operators.
End of explanation
"""


def project(df, columns):
    return df[columns] #SOLUTION

project(salty_sailors, ["sname", "age"])

_ = ok.grade('qproject')
_ = ok.backup()

"""
Explanation: Question 1: Projection
Our arguments are a dataframe and a list of columns to select. This should be a simple one :)
End of explanation
"""


def select(df, condition):
    return df[df.apply(condition, axis=1)] #SOLUTION

select(young_sailors, lambda x: x["rating"] > 6)

_ = ok.grade('qselect')
_ = ok.backup()

"""
Explanation: Question 2: Selection
For selecton, our arguments are a dataframe and a function which determines which rows we select. For instance,
good_sailors = select(young_sailors, lambda x: x["rating"] &gt; 6)
End of explanation
"""


def union(df1, df2):
    return pd.concat([df1, df2]).drop_duplicates() #SOLUTION

union(young_sailors, salty_sailors)

_ = ok.grade('qunion')
_ = ok.backup()

"""
Explanation: Question 3: Union
This is a binary operator, so we pass in two dataframes as our arguments. You can assume that the two dataframes are union compatible - that is, that they have the same number of columns, and their columns have the same types.
End of explanation
"""


def intersection(df1, df2):
    return pd.merge(df1, df2) #SOLUTION

intersection(young_sailors, salty_sailors)

_ = ok.grade('qintersection')
_ = ok.backup()

"""
Explanation: Question 4: Intersection
Similar to Union, this is also a binary operator.
End of explanation
"""


def difference(df1, df2):
    return df1.where(df1.apply(lambda x: ~x.isin(df2[x.name]))).dropna()

difference(young_sailors, salty_sailors)

_ = ok.grade('qdifference')
_ = ok.backup()

"""
Explanation: Question 5: Set-difference
This one is a bit harder. You might just want to convert the rows of the dataframes to tuple, if you're having trouble.
End of explanation
"""


def cross_product(df1, df2):
    # add a column "tmp-key" of zeros to df1 and df2 
    df1 = pd.concat([df1, pd.Series(0, index=df1.index, name="tmp-key")], axis=1)
    df2 = pd.concat([df2, pd.Series(0, index=df2.index, name="tmp-key")], axis=1)
    # use Pandas merge functionality along with drop 
    # to compute outer product and remove extra column
    return (pd
            .merge(df1, df2, on="tmp-key")
            .drop(["tmp-key"], axis=1)) #SOLUTION

cross_product(young_sailors, salty_sailors)

_ = ok.grade('qcross_product')
_ = ok.backup()

"""
Explanation: Question 6: Cross-product
This one is also tricky, so we've provided some help for you. Think about how the new key column could be used...
End of explanation
"""


def theta_join(df1, df2, condition):
    return select(cross_product(df1, df2), condition)

theta_join(young_sailors, salty_sailors, lambda x: x["age_x"] > x["age_y"])

_ = ok.grade('qtheta_join')
_ = ok.backup()

"""
Explanation: Question 7: Theta-Join
Can you do this by using two other relational operators?
End of explanation
"""


def natural_join(df1, df2, attr):
    return select(cross_product(df1, df2), lambda x: x[attr+"_x"] == x[attr+"_y"])

all_sailors = union(young_sailors, salty_sailors)
sailor_reservtions = natural_join(all_sailors, reservations, "sid")
sailors_and_boats = natural_join(sailor_reservtions, boats, "bid")
project(sailors_and_boats, ["sname", "bname", "day"])

_ = ok.grade('qnatural_join')
_ = ok.backup()

"""
Explanation: Question 8: Natural Join
Similar to above, try to implement this using two relational operators.
End of explanation
"""


i_finished_the_lab = False

_ = ok.grade('qcompleted')
_ = ok.backup()

_ = ok.submit()

"""
Explanation: Submitting your assignment
If you made a good-faith effort to complete the lab, change i_finished_the_lab to True in the cell below.  In any case, run the cells below to submit the lab.
End of explanation
"""


