import hail as hl

hl.utils.get_movie_lens('data/')
users = hl.read_table('data/users.ht')

users.filter(users.occupation == 'programmer').count()

"""
Explanation: Filtering and Annotation Tutorial
Filter
You can filter the rows of a table with Table.filter.   This returns a table of those rows for which the expression evaluates to True.
End of explanation
"""


users.aggregate(hl.agg.filter(users.occupation == 'programmer', hl.agg.count()))

users.aggregate(hl.agg.counter(users.occupation == 'programmer'))[True]

"""
Explanation: We can also express this query in multiple ways using aggregations:
End of explanation
"""


missing_occupations = hl.set(['other', 'none'])

t = users.annotate(
    cleaned_occupation = hl.if_else(missing_occupations.contains(users.occupation),
                                    hl.missing('str'),
                                    users.occupation))
t.show()

"""
Explanation: Annotate
You can add new fields to a table with annotate.  As an example, let's create a new column called cleaned_occupation that replaces missing entries in the occupation field labeled as 'other' with 'none.'
End of explanation
"""


users.show()

"""
Explanation: Compare this to what we had before:
End of explanation
"""


users.describe()

"""
Explanation: Note: annotate is functional: it doesn't mutate users, but returns a new table.  This is also true of filter.  In fact, all operations in Hail are functional.
End of explanation
"""


users.select(users.sex, users.occupation).show()

"""
Explanation: Select and Transmute
There are two other annotate methods: select and transmute.  select allows you to create new tables from old ones by selecting existing fields, or creating new ones.
First, let's extract the sex and occupation fields:
End of explanation
"""


mean_age = round(users.aggregate(hl.agg.stats(users.age)).mean)
users.select(users.sex, users.occupation, mean_shifted_age = users.age - mean_age).show()

"""
Explanation: We can also create a new field that stores the age relative to the average. Note that new fields must be assigned a name (in this case mean_shifted_age):
End of explanation
"""


missing_occupations = hl.set(['other', 'none'])

t = users.select(
    cleaned_occupation = hl.if_else(missing_occupations.contains(users.occupation),
                                    hl.missing('str'),
                                    users.occupation))
t.show()

missing_occupations = hl.set(['other', 'none'])

t = users.transmute(
    cleaned_occupation = hl.if_else(missing_occupations.contains(users.occupation),
                                    hl.missing('str'),
                                    users.occupation))
t.show()

"""
Explanation: transmute replaces any fields mentioned on the right-hand side with the new fields, but leaves unmentioned fields unchanged.  transmute is useful for transforming data into a new form. Compare the following two snippts of code. The second is identical to the first with transmute replacing select.
End of explanation
"""


t = users.annotate_globals(cohort = 5, cloudable = hl.set(['sample1', 'sample10', 'sample15']))
t.describe()

t.cloudable

hl.eval(t.cloudable)

"""
Explanation: Global Fields
Finally, you can add global fields with annotate_globals.  Globals are useful for storing metadata about a dataset or storing small data structures like sets and maps.
End of explanation
"""


