'''Assignment 2 - Introduction to NLTK
In part 1 of this assignment you will use nltk to explore the Herman Melville novel Moby Dick. Then in
part 2 you will create a spelling recommender function that uses nltk to find words similar to the
misspelling.'''

'''Part 1 - Analyzing Moby Dick'''

import nltk
import pandas as pd
import numpy as np

nltk.download('punkt')
# If you would like to work with the raw text you can use 'moby_raw'
with open('moby.txt', 'r') as f:
    moby_raw = f.read()

# If you would like to work with the novel in nltk.Text format you can use 'text1'
moby_tokens = nltk.word_tokenize(moby_raw)
text1 = nltk.Text(moby_tokens)

#-----------------------------------------------------------------------
'''Example 1
How many tokens (words and punctuation symbols) are in text1?

This function should return an integer.'''

#---------- ANSWER CODE ----------
def example_one():

    return len(nltk.word_tokenize(moby_raw)) # or alternatively len(text1)

example_one()

#---------- ANSWER ----------
255038

#-----------------------------------------------------------------------
'''Example 2
How many unique tokens (unique words and punctuation) does text1 have?

This function should return an integer.'''

#---------- ANSWER CODE ----------
def example_two():

    return len(set(nltk.word_tokenize(moby_raw))) # or alternatively len(set(text1))

example_two()

#---------- ANSWER ----------
20742

#-----------------------------------------------------------------------
'''Example 3
After lemmatizing the verbs, how many unique tokens does text1 have?

This function should return an integer.'''

#---------- ANSWER CODE ----------
from nltk.stem import WordNetLemmatizer
nltk.download('wordnet')

def example_three():

    lemmatizer = WordNetLemmatizer()
    lemmatized = [lemmatizer.lemmatize(w,'v') for w in text1]

    return len(set(lemmatized))

example_three()

#---------- ANSWER ----------
16887

#-----------------------------------------------------------------------
'''Question 1
What is the lexical diversity of the given text input? (i.e. ratio of unique tokens to the total number of
tokens)

This function should return a float.'''

#---------- ANSWER CODE ----------
def answer_one():

    return len(set(moby_tokens))/len(moby_tokens)

answer_one()

#---------- ANSWER ----------
0.08132905684643073

#-----------------------------------------------------------------------
'''Question 2
What percentage of tokens is 'whale'or 'Whale'?

This function should return a float.'''

#---------- ANSWER CODE ----------
def answer_two():

    dist = nltk.FreqDist(text1)
    return (dist['whale']+ dist['Whale']) / len(moby_tokens)*100# Your answer here

answer_two()

#---------- ANSWER ----------
0.41248755087477157

#-----------------------------------------------------------------------
'''Question 3
What are the 20 most frequently occurring (unique) tokens in the text? What is their frequency?

This function should return a list of 20 tuples where each tuple is of the form (token, frequency).
The list should be sorted in descending order of frequency'''

#---------- ANSWER CODE ----------
def answer_three():
    from collections import Counter
    dist = nltk.FreqDist(text1)
    return Counter(dist).most_common(20)
#     sorted_x = sorted(dist.items(), key=lambda kv: kv[1])
#     return sorted_x[::-1][:20]# Your answer here

answer_three()

#---------- ANSWER ----------
'''
[(',', 19204),
 ('the', 13715),
 ('.', 7306),
 ('of', 6513),
 ('and', 6010),
 ('a', 4545),
 ('to', 4515),
 (';', 4173),
 ('in', 3908),
 ('that', 2978),
 ('his', 2459),
 ('it', 2196),
 ('I', 2113),
 ('!', 1767),
 ('is', 1722),
 ('--', 1713),
 ('with', 1659),
 ('he', 1658),
 ('was', 1639),
 ('as', 1620)]'''

#-----------------------------------------------------------------------
'''Question 4
What tokens have a length of greater than 5 and frequency of more than 150?

This function should return an alphabetically sorted list of the tokens that match the above constraints.
To sort your list, use sorted()'''

#---------- ANSWER CODE ----------
def answer_four():
    dist = nltk.FreqDist(text1)
    vocab1 = dist.keys()
    freq = [w for w in vocab1 if len(w) > 5 and dist[w] > 150]

    return sorted(freq)# Your answer here

answer_four()

#---------- ANSWER ----------
'''
['Captain',
 'Pequod',
 'Queequeg',
 'Starbuck',
 'almost',
 'before',
 'himself',
 'little',
 'seemed',
 'should',
 'though',
 'through',
 'whales',
 'without']'''

#-----------------------------------------------------------------------
'''Question 5
Find the longest word in text1 and that word's length.

This function should return a tuple (longest_word, length)'''

#---------- ANSWER CODE ----------
def answer_five():
    mt = np.array(list(set(moby_tokens)))
    lens = np.array([len(w) for w in mt])

    return mt[lens == np.max(lens)][0],np.max(lens)#lens[lens == np.max(lens)]# Your answer here

answer_five()

#---------- ANSWER ----------
("twelve-o'clock-at-night", 23)

#-----------------------------------------------------------------------
'''Question 6
What unique words have a frequency of more than 2000? What is their frequency?

"Hint: you may want to use isalpha() to check if the token is a word and not punctuation."

This function should return a list of tuples of the form (frequency, word) sorted in descending order of
frequency.'''

#---------- ANSWER CODE ----------
def answer_six():
    unique = answer_three()
    answer = [(f,w) for w,f in unique if w.isalpha() and f>2000]
    return answer

answer_six()

#---------- ANSWER ----------
'''
[(13715, 'the'),
 (6513, 'of'),
 (6010, 'and'),
 (4545, 'a'),
 (4515, 'to'),
 (3908, 'in'),
 (2978, 'that'),
 (2459, 'his'),
 (2196, 'it'),
 (2113, 'I')]'''

#-----------------------------------------------------------------------
'''Question 7
What is the average number of tokens per sentence?

This function should return a float.'''

#---------- ANSWER CODE ----------
def answer_seven():
    sentences = nltk.sent_tokenize(moby_raw)
    sentence_token = np.array([len(nltk.word_tokenize(s)) for s in sentences])

    return sentence_token.mean()# Your answer here

answer_seven()

#---------- ANSWER ----------
25.886926512383273

#-----------------------------------------------------------------------
'''Question 8
What are the 5 most frequent parts of speech in this text? What is their frequency?

This function should return a list of tuples of the form (part_of_speech, frequency) sorted in descending
order of frequency.'''

#---------- ANSWER CODE ----------
nltk.download('averaged_perceptron_tagger')
def answer_eight():
    from collections import Counter
    pos_tag = nltk.pos_tag(text1) #moby_tokens)
    tag = [t for p,t in pos_tag ]
    dist = nltk.FreqDist(tag)

    return Counter(dist).most_common(5)# Your answer here

answer_eight()

#---------- ANSWER ----------
[('NN', 32729), ('IN', 28663), ('DT', 25879), (',', 19204), ('JJ', 17613)]

#-----------------------------------------------------------------------
'''Part 2 - Spelling Recommender
For this part of the assignment you will create three different spelling recommenders, that each take a
list of misspelled words and recommends a correctly spelled word for every word in the list.
For every misspelled word, the recommender should find find the word in correct_spellings that has the
shortest distance*, and starts with the same letter as the misspelled word, and return that word as a
recommendation.
*Each of the three different recommenders will use a different distance measure (outlined below).
Each of the recommenders should provide recommendations for the three default words provided:
['cormulent', 'incendenece', 'validrate'].'''

from nltk.corpus import words
nltk.download('words')
correct_spellings = words.words()

#-----------------------------------------------------------------------
'''Question 9
For this recommender, your function should provide recommendations for the three default words provided above using the following distance
metric:

Jaccard distance on the trigrams of the two words.

This function should return a list of length three:
['cormulent_reccomendation', 'incendenece_reccomendation', 'validrate_reccomendation'].'''

#---------- ANSWER CODE ----------
def answer_nine(entries=['cormulent', 'incendenece', 'validrate']):

    dic = {}
    for entrie in entries:

        jd = [( word,
                nltk.jaccard_distance(
                    set(nltk.ngrams(entrie,n=3)),
                    set(nltk.ngrams(word,n=3))
                )
              ) for word in correct_spellings if word[0] == entrie[0]
            ]
        dic[entrie] = pd.DataFrame(jd,columns =['word','jd'])

    answer = [dic[k][dic[k].jd==dic[k].jd.min()].word.values[0] for k in entries]

    return answer # Your answer here

answer_nine()

#---------- ANSWER ----------
['corpulent', 'indecence', 'validate']

#-----------------------------------------------------------------------
'''Question 10
For this recommender, your function should provide recommendations for the three default words provided
above using the following distance metric:

Jaccard distance on the 4-grams of the two words.

This function should return a list of length three:
['cormulent_reccomendation', 'incendenece_reccomendation', 'validrate_reccomendation'].'''

#---------- ANSWER CODE ----------
def answer_ten(entries=['cormulent', 'incendenece', 'validrate']):

    dic = {}
    for entrie in entries:

        jd = [( word,
                nltk.jaccard_distance(
                    set(nltk.ngrams(entrie,n=4)),
                    set(nltk.ngrams(word,n=4))
                )
              ) for word in correct_spellings if word[0] == entrie[0]
            ]
        dic[entrie] = pd.DataFrame(jd,columns =['word','jd'])

    answer = [dic[k][dic[k].jd==dic[k].jd.min()].word.values[0] for k in entries]

    return answer # Your answer here

answer_ten()

#---------- ANSWER ----------
['cormus', 'incendiary', 'valid']

#-----------------------------------------------------------------------
'''Question 11
For this recommender, your function should provide recommendations for the three default words provided
above using the following distance metric:

Edit distance on the two words with transpositions.

This function should return a list of length three:
['cormulent_reccomendation', 'incendenece_reccomendation', 'validrate_reccomendation'].'''

#---------- ANSWER CODE ----------
def answer_eleven(entries=['cormulent', 'incendenece', 'validrate']):

    dic = {}

    for entrie in entries:
        jd = [( word,
                nltk.edit_distance(entrie, word)
              ) for word in correct_spellings if word[0] == entrie[0]
            ]
        dic[entrie] = pd.DataFrame(jd,columns =['word','jd'])

    answer = [dic[k][dic[k].jd==dic[k].jd.min()].word.values[0] for k in entries]

    return answer # Your answer here

answer_eleven()
#---------- ANSWER ----------
['corpulent', 'intendence', 'validate']
#-----------------------------------------------------------------------


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:light
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.3'
#       jupytext_version: 0.8.2
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
#   language_info:
#     codemirror_mode:
#       name: ipython
#       version: 3
#     file_extension: .py
#     mimetype: text/x-python
#     name: python
#     nbconvert_exporter: python
#     pygments_lexer: ipython3
#     version: 3.6.5
# ---

# %load_ext autoreload
# %autoreload 2

import sys

sys.path.append("../..")

from optimus import Optimus

op = Optimus(master='local')

# +
import pandas as pd
from pyspark.sql.types import *
from datetime import date, datetime

cols = [
        ("names", "str"),
        ("height(ft)", ShortType()),
        ("function", "str"),
        ("rank", ByteType()),
        ("age", "int"),
        ("weight(t)", "float"),
        "japanese name",
        "last position seen",
        "date arrival",
        "last date seen",
        ("attributes", ArrayType(FloatType())),
        ("DateType", DateType()),
        ("Tiemstamp", TimestampType()),
        ("Cybertronian", BooleanType()),
        ("function(binary)", BinaryType()),
        ("NullType", NullType())

    ]

rows = [
        ("Optim'us", 28, "Leader", 10, 5000000, 4.30, ["Inochi", "Convoy"], "19.442735,-99.201111", "1980/04/10",
         "2016/09/10", [8.5344, 4300.0], date(2016, 9, 10), datetime(2014, 6, 24), True, bytearray("Leader", "utf-8"),
         None),
        ("bumbl#ebéé  ", 17, "Espionage", 7, 5000000, 2.0, ["Bumble", "Goldback"], "10.642707,-71.612534", "1980/04/10",
         "2015/08/10", [5.334, 2000.0], date(2015, 8, 10), datetime(2014, 6, 24), True, bytearray("Espionage", "utf-8"),
         None),
        ("ironhide&", 26, "Security", 7, 5000000, 4.0, ["Roadbuster"], "37.789563,-122.400356", "1980/04/10",
         "2014/07/10", [7.9248, 4000.0], date(2014, 6, 24), datetime(2014, 6, 24), True, bytearray("Security", "utf-8"),
         None),
        ("Jazz", 13, "First Lieutenant", 8, 5000000, 1.80, ["Meister"], "33.670666,-117.841553", "1980/04/10",
         "2013/06/10", [3.9624, 1800.0], date(2013, 6, 24), datetime(2014, 6, 24), True,
         bytearray("First Lieutenant", "utf-8"), None),
        ("Megatron", None, "None", 10, 5000000, 5.70, ["Megatron"], None, "1980/04/10", "2012/05/10", [None, 5700.0],
         date(2012, 5, 10), datetime(2014, 6, 24), True, bytearray("None", "utf-8"), None),
        ("Metroplex_)^$", 300, "Battle Station", 8, 5000000, None, ["Metroflex"], None, "1980/04/10", "2011/04/10",
         [91.44, None], date(2011, 4, 10), datetime(2014, 6, 24), True, bytearray("Battle Station", "utf-8"), None),

    ]
df = op.create.df(cols ,rows)
df.table()
# -

from optimus.helpers.test import Test
from pyspark.ml.linalg import Vectors

# +
## Optimus Test
# -

t = Test(op, None, "Optimus", imports=["import datetime",
                                "from pyspark.sql import functions as F"])

# +
one_column = {"rows":["Argenis", "Favio", "Matthew"], "cols":["name"]}
plain = {"rows":[("BOB", 1),("JoSe", 2)],"cols":["name","age"]}
plain_infer_false = {"rows":[("BOB", 1),("JoSe", 2)],"cols":["name","age"],"infer_schema":False}
with_data_types = {"rows":[("BOB", 1),("JoSe", 2)],"cols":[("name", StringType(), True),("age", IntegerType(), False)]}
nullable = {"rows":[("BOB", 1),("JoSe", 2)],"cols":[("name", StringType()),("age", IntegerType())]}

df1 = op.create.df(**one_column)
df2 = op.create.df(**plain)
df3 = op.create.df(**plain_infer_false)
df4 = op.create.df(**with_data_types)
df5 = op.create.df(**nullable)

t.run(

    t.create(df1, None, "one_column", "df", **one_column),
    t.create(df2, None, "plain", "df", **plain),
    t.create(df3, None, "plain_infer_false", "df", **plain_infer_false),
    t.create(df4, None, "with_data_types", "df", **with_data_types),
    t.create(df5, None, "nullable", "df", **nullable),

)

# +
## Columns Test
# -

t = Test(op, df, "df_cols", imports=["from pyspark.ml.linalg import Vectors, VectorUDT, DenseVector",
                                        "import numpy as np",
                                        "nan = np.nan",
                                        "import datetime",
                                        "from pyspark.sql import functions as F"])

# +
from pyspark.sql import functions as F

def func(col_name, attrs):
    return F.col(col_name) * 2

numeric_col = "height(ft)"
numeric_col_B = "rank"
numeric_col_C = "rank"
string_col = "function"
data_col = "date arrival"
data_col_B = "last date seen"
new_col = "new col"
array_col = "attributes"

t.run(

    t.create(None, "cols.min", None, "json", numeric_col),
    t.create(None, "cols.min", "all_columns", "json", "*"),

    t.create(None, "cols.max", None, "json", numeric_col),
    t.create(None, "cols.max", "all_columns", "json", "*"),

    t.create(None, "cols.range", None, "json", numeric_col),
    t.create(None, "cols.range", "all_columns", "json", "*"),

    t.create(None, "cols.median", None, "json", numeric_col),
    t.create(None, "cols.median", "all_columns", "json", "*"),

    t.create(None, "cols.percentile", None, "json", numeric_col, [0.05, 0.25], 1),
    t.create(None, "cols.percentile", "all_columns", "json", "*", [0.05, 0.25], 1),

    t.create(None, "cols.mad", None, "json", numeric_col),
    t.create(None, "cols.mad", "all_columns", "json", "*"),

    t.create(None, "cols.std", None, "json", numeric_col),
    t.create(None, "cols.std", "all_columns", "json", "*"),

    t.create(None, "cols.kurt", None, "json", numeric_col),
    t.create(None, "cols.kurt", "all_columns", "json", "*"),

    t.create(None, "cols.mean", None, "json", numeric_col),
    t.create(None, "cols.mean", "all_columns", "json", "*"),

    t.create(None, "cols.skewness", None, "json", numeric_col),
    t.create(None, "cols.skewness", "all_columns", "json", "*"),

    t.create(None, "cols.sum", None, "json", numeric_col),
    t.create(None, "cols.sum", "all_columns", "json", "*"),

    t.create(None, "cols.variance", None, "json", numeric_col),
    t.create(None, "cols.variance", "all_columns", "json", "*"),

    t.create(None, "cols.abs", None, "df", numeric_col),
    t.create(None, "cols.abs", "all_columns", "df", "*"),

    t.create(None, "cols.mode", None, "json", numeric_col),
    t.create(None, "cols.mode", "all_columns", "json", "*"),

    t.create(None, "cols.count", None, "json"),

    t.create(None, "cols.count_na", None, "json", numeric_col),
    t.create(None, "cols.count_na", "all_columns", "json", "*"),

    t.create(None, "cols.count_zeros", None, "json", numeric_col),
    t.create(None, "cols.count_zeros", "all_columns", "json", "*"),

    t.create(None, "cols.count_uniques", None, "json", numeric_col),
    t.create(None, "cols.count_uniques", "all_columns", "json", "*"),

    t.create(None, "cols.unique", None, "df", numeric_col),
    t.create(None, "cols.unique", "all_columns", "df", "*"),

    t.create(None, "cols.add", None, "df", [numeric_col, numeric_col_B]),
    t.create(None, "cols.add", "all_columns", "df", "*"),

    t.create(None, "cols.sub", None, "df", [numeric_col, numeric_col_B]),
    t.create(None, "cols.sub", "all_columns", "df", "*"),

    t.create(None, "cols.mul", None, "df", [numeric_col, numeric_col_B]),
    t.create(None, "cols.mul", "all_columns", "df", "*"),

    t.create(None, "cols.div", None, "df", [numeric_col, numeric_col_B]),
    t.create(None, "cols.div", "all_columns", "df", "*"),

    t.create(None, "cols.z_score", None, "df", numeric_col),
    t.create(None, "cols.z_score", "all_columns", "df", "*"),

    t.create(None, "cols.iqr", None, "json", numeric_col),
    t.create(None, "cols.iqr", "all_columns", "json", "*"),

    t.create(None, "cols.lower", None, "df", numeric_col),
    t.create(None, "cols.lower", "all_columns", "df", "*"),

    t.create(None, "cols.upper", None, "df", numeric_col),
    t.create(None, "cols.upper", "all_columns", "df", "*"),

    t.create(None, "cols.trim", None, "df", numeric_col),

    t.create(None, "cols.trim", "all_columns", "df", "*"),

    t.create(None, "cols.reverse", None, "df", numeric_col),
    t.create(None, "cols.reverse", "all_columns", "df", "*"),

    t.create(None, "cols.remove_accents", None, "df", numeric_col),
    t.create(None, "cols.remove_accents", "all_columns", "df", "*"),

    t.create(None, "cols.remove_special_chars", None, "df", numeric_col),
    t.create(None, "cols.remove_special_chars", "all_columns", "df", "*"),

    t.create(None, "cols.remove_white_spaces", None, "df", numeric_col),
    t.create(None, "cols.remove_white_spaces", "all_columns", "df", "*"),

    t.create(None, "cols.date_transform", None, "df", data_col, "yyyy/MM/dd", "dd-MM-YYYY"),
    t.create(None, "cols.date_transform", "all_columns", "df", [data_col, data_col_B], "yyyy/MM/dd", "dd-MM-YYYY"),

    t.create(None, "cols.years_between", None, "df", data_col, "yyyyMMdd"),
    t.create(None, "cols.years_between", "multiple_columns", "df", [data_col, data_col_B], "yyyyMMdd"),

    # ---

    t.create(None, "cols.impute", None, "df", numeric_col_B),
    t.create(None, "cols.impute", "all_columns", "df", "*"),

    t.create(None, "cols.hist", None, "json", numeric_col_B, 4),
    #t.create(None,"cols.hist","all_columns","df","*",4),

    t.create(None, "cols.frequency", None, "json", numeric_col_B, 4),
    t.create(None, "cols.frequency", "all_columns", "json", "*", 4),

    t.create(None, "cols.schema_dtype", None, "json", numeric_col_B),
    #t.create(None, "cols.schema_dtype", "all_columns", "json", "*"),

    t.create(None, "cols.dtypes", None, "json", numeric_col_B),
    t.create(None, "cols.dtypes", "all_columns", "json", "*"),

    t.create(None, "cols.select_by_dtypes", "str", "df", "str"),
    t.create(None, "cols.select_by_dtypes", "int", "df", "int"),
    t.create(None, "cols.select_by_dtypes", "float", "df", "float"),
    t.create(None, "cols.select_by_dtypes", "array", "df", "array"),

    t.create(None, "cols.names", None, "json"),

    t.create(None, "cols.qcut", None, "df", numeric_col_B, 4),
    t.create(None, "cols.qcut", "all_columns", "df", "*", 4),

    t.create(None, "cols.clip", None, "df", numeric_col_B, 3, 5),
    t.create(None, "cols.clip", "all_columns", "df", "*", 3, 5),

    t.create(None, "cols.replace", None, "df", string_col, [("Security", "Leader")], "Match"),
    t.create(None, "cols.replace", "all_columns", "df", "*", [("Jazz", "Leader")], "Match"),

    t.create(None, "cols.apply_expr", None, "df", numeric_col_B, func),
    t.create(None, "cols.apply_expr", "all_columns", "df", [numeric_col_B,numeric_col_C], func),

    t.create(None, "cols.append", "number", "df", new_col, 1),

    #t.create(None, "cols.append", "advance", "df", [("new_col_4", "test"),
    #                                                ("new_col_5", df[numeric_col_B] * 2),
    #                                                ("new_col_6", [1, 2, 3])
    #                                                ]),

    t.create(None, "cols.rename", None, "df", numeric_col_B, numeric_col_B + "(old)"),
    t.create(None, "cols.rename", "list", "df",
             [numeric_col, numeric_col + "(tons)", numeric_col_B, numeric_col_B + "(old)"]),
    t.create(None, "cols.rename", "function", "df", str.upper),

    t.create(None, "cols.drop", None, "df", numeric_col_B),

    t.create(None, "cols.cast", None, "df", string_col, "string"),
    t.create(None, "cols.cast", "all_columns", "df", "*", "string"),
    t.create(None, "cols.cast", "vector", "df", array_col, Vectors),

    t.create(None, "cols.keep", None, "df", numeric_col_B),

    t.create(None, "cols.move", None, "df", numeric_col_B, "after", array_col),

    t.create(None, "cols.select", None, "df", 0, numeric_col),

    t.create(None, "cols.select", "regex", "df", "n.*", regex=True),

    t.create(None, "cols.sort", None, "df"),
    t.create(None, "cols.sort", "desc", "df", "desc"),
    t.create(None, "cols.sort", "asc", "df", "asc"),

    t.create(None, "cols.fill_na", None, "df", numeric_col, "N/A"),
    t.create(None, "cols.fill_na", "all_columns", "df", "*", "N/A"),

    t.create(None, "cols.nest", None, "df", [numeric_col, numeric_col_B], new_col, separator=" "),
    #t.create(None, "cols.nest", "mix", "df", [F.col(numeric_col), F.col(numeric_col_B)], "E", separator="--"),

    #t.create(None, "cols.nest", "vector_all_columns", "df", [numeric_col, numeric_col_B], new_col, shape="vector"),
    t.create(None, "cols.nest", "vector", "df", [numeric_col_C, numeric_col_B], new_col, shape="vector"),

    #t.create(None, "cols.nest", "array_all_columns", "df", "*", new_col, shape="array"),
    t.create(None, "cols.nest", "array", "df", [numeric_col, numeric_col_B,numeric_col_C], new_col, shape="array"),

    t.create(None, "cols.unnest", "array_all_columns", "df", array_col, "-", index=1),
    t.create(None, "cols.unnest", "array", "df", array_col),
    t.create(None, "cols.unnest", "array_all_columns", "df", array_col),

    t.create(None, "cols.is_na", "all_columns", "df", "*"),
    t.create(None, "cols.is_na", None, "df", numeric_col),

)
# -

t = Test(op,df, "df_rows", imports=["from pyspark.ml.linalg import Vectors, VectorUDT, DenseVector",
                                         "import numpy as np",
                                        "nan = np.nan",
                                        "import datetime",
                                        "from pyspark.sql import functions as F",
                                        "from optimus.functions import abstract_udf as audf"])

rows = [
        ("Optim'us", 28, "Leader", 10, 5000000, 4.30, ["Inochi", "Convoy"], "19.442735,-99.201111", "1980/04/10",
         "2016/09/10", [8.5344, 4300.0], date(2016, 9, 10), datetime(2014, 6, 24), True, bytearray("Leader", "utf-8"),
         None)
]

# +
from pyspark.sql import functions as F
from optimus.functions import abstract_udf as audf

def func_data_type(value, attr):
    return value > 1

t.run(

    t.create(None, "rows.append", None, "df", rows),
    #t.create(None, "rows.select", None, "df", F.col("rank") == 7),
    t.create(None, "rows.select_by_dtypes", "integer", "df", "height(ft)", "integer"),
    t.create(None, "rows.select_by_dtypes", "float", "df", "weight(t)", "float"),

    t.create(None, "rows.drop_by_dtypes", "integer", "df", "height(ft)", "integer"),
    t.create(None, "rows.drop_by_dtypes", "float", "df", "weight(t)", "float"),

    #t.create(None, "rows.drop", None, "df", (F.col("rank") == 10) | (F.col("rank") == 7)),
    #t.create(None, "rows.drop", "audf", "df", (audf("rank", func_data_type, "boolean"))),

    t.create(None, "rows.sort", None, "df","rank"),
    t.create(None, "rows.sort", "desc", "df", "rank", "desc"),
    t.create(None, "rows.sort", "asc", "df", "rank", "asc"),

    #t.create(None, "rows.is_in", None, "df", ("rank", 2)),
)

# -


# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import print_function

import unittest
import paddle
import numpy as np
import paddle.fluid.core as core
from op_test import OpTest
import paddle.fluid as fluid
from paddle.fluid import Program, program_guard

class TestIndexSelectOp(OpTest):
    def setUp(self):
        self.op_type = "index_select"
        self.init_dtype_type()
        index_np = np.random.randint(
            low=0, high=self.x_shape[self.dim], size=self.index_size)
        x_np = np.random.random(self.x_shape).astype(self.x_type)
        self.inputs = {'X': x_np, 'Index': index_np}
        self.attrs = {'dim': self.dim}
        outer_loop = np.prod(self.x_shape[:self.dim])
        x_reshape = [outer_loop] + list(self.x_shape[self.dim:])
        x_np_reshape = np.reshape(x_np, tuple(x_reshape))
        out_list = []
        for i in range(outer_loop):
            for j in range(self.index_size):
                out_list.append(x_np_reshape[i, index_np[j]])
        self.out_shape = list(self.x_shape)
        self.out_shape[self.dim] = self.index_size
        self.out_shape = tuple(self.out_shape)

        out = np.reshape(out_list, self.out_shape)
        self.outputs = {'Out': out}

    def init_dtype_type(self):
        self.dim = 1
        self.x_type = np.float64
        self.index_type = np.int64
        self.x_shape = (100, 4, 5)
        self.index_size = 100

    def test_check_output(self):
        self.check_output()

    def test_check_grad_normal(self):
        self.check_grad(['X'], 'Out')

class TestIndexSelectOpCase2(TestIndexSelectOp):
    def init_dtype_type(self):
        self.x_type = np.float32
        self.index_type = np.int32
        self.dim = -2
        self.x_shape = (10, 10, 4, 10)
        self.index_size = 10

class TestIndexSelectOpCaseSingleThread(TestIndexSelectOp):
    def init_dtype_type(self):
        if fluid.is_compiled_with_cuda():
            fluid.set_flags({'FLAGS_cudnn_deterministic': True})
        self.x_type = np.float32
        self.index_type = np.int32
        self.dim = -2
        self.x_shape = (10, 10, 4, 10)
        self.index_size = 10

class TestIndexSelectAPI(unittest.TestCase):
    def input_data(self):
        self.data_x = np.array([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0],
                                [9.0, 10.0, 11.0, 12.0]])
        self.data_index = np.array([0, 1, 1]).astype('int32')

    def test_index_select_api(self):
        self.input_data()

        # case 1:
        with program_guard(Program(), Program()):
            x = fluid.layers.data(name='x', shape=[-1, 4])
            index = fluid.layers.data(
                name='index', shape=[3], dtype='int32', append_batch_size=False)
            z = paddle.index_select(x, index, axis=1)
            exe = fluid.Executor(fluid.CPUPlace())
            res, = exe.run(feed={'x': self.data_x,
                                 'index': self.data_index},
                           fetch_list=[z.name],
                           return_numpy=False)
        expect_out = np.array([[1.0, 2.0, 2.0], [5.0, 6.0, 6.0],
                               [9.0, 10.0, 10.0]])
        self.assertTrue(np.allclose(expect_out, np.array(res)))

        # case 2:
        with program_guard(Program(), Program()):
            x = fluid.layers.data(name='x', shape=[-1, 4])
            index = fluid.layers.data(
                name='index', shape=[3], dtype='int32', append_batch_size=False)
            z = paddle.index_select(x, index)
            exe = fluid.Executor(fluid.CPUPlace())
            res, = exe.run(feed={'x': self.data_x,
                                 'index': self.data_index},
                           fetch_list=[z.name],
                           return_numpy=False)
        expect_out = np.array(
            [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [5.0, 6.0, 7.0, 8.0]])
        self.assertTrue(np.allclose(expect_out, np.array(res)))

    def test_dygraph_api(self):
        self.input_data()
        # case 1:
        with fluid.dygraph.guard():
            x = fluid.dygraph.to_variable(self.data_x)
            index = fluid.dygraph.to_variable(self.data_index)
            z = paddle.index_select(x, index)
            np_z = z.numpy()
        expect_out = np.array(
            [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [5.0, 6.0, 7.0, 8.0]])
        self.assertTrue(np.allclose(expect_out, np_z))

        # case 2:
        with fluid.dygraph.guard():
            x = fluid.dygraph.to_variable(self.data_x)
            index = fluid.dygraph.to_variable(self.data_index)
            z = paddle.index_select(x, index, axis=1)
            np_z = z.numpy()
        expect_out = np.array([[1.0, 2.0, 2.0], [5.0, 6.0, 6.0],
                               [9.0, 10.0, 10.0]])
        self.assertTrue(np.allclose(expect_out, np_z))

if __name__ == '__main__':
    unittest.main()


import os
import random
import numpy as np
from tensorflow.random import set_seed
SEED = 42

def setup_seed():
    os.environ["TF_DETERMINISTIC_OPS"] = "1"
    os.environ["PYTHONHASHSEED"] = str(SEED)
    random.seed(SEED)
    np.random.seed(SEED)
    set_seed(SEED)
    return


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed May 30 05:57:34 2018

@author: prmiles
"""

from pymcmcstat.plotting import utilities
import unittest
from mock import patch
import numpy as np
import math

# --------------------------
class GenerateSubplotGrid(unittest.TestCase):
    def test_generate_subplot_grid(self):
        nparam = 5
        ns1, ns2 = utilities.generate_subplot_grid(nparam = nparam)
        self.assertEqual(ns1, math.ceil(math.sqrt(nparam)), msg = 'Expect 3')
        self.assertEqual(ns2, round(math.sqrt(nparam)), msg = 'Expect 2')

    def test_generate_subplot_grid_1(self):
        nparam = 1
        ns1, ns2 = utilities.generate_subplot_grid(nparam = nparam)
        self.assertEqual(ns1, math.ceil(math.sqrt(nparam)), msg = 'Expect 1')
        self.assertEqual(ns2, round(math.sqrt(nparam)), msg = 'Expect 1')

# --------------------------
class GenerateNames(unittest.TestCase):

    def test_default_names(self):
        nparam = 8
        names = utilities.generate_names(nparam = nparam, names = None)
        self.assertEqual(len(names), nparam, msg = 'Length of names should match number of parameters')
        for ii in range(nparam):
            self.assertEqual(names[ii], str('$p_{{{}}}$'.format(ii)))

    def test_names_partial(self):
        nparam = 8
        names = ['hi']
        names = utilities.generate_names(nparam = nparam, names = names)
        self.assertEqual(names[0], 'hi', msg = 'First name is hi')
        for ii in range(1, nparam):
            self.assertEqual(names[ii], str('$p_{{{}}}$'.format(ii)))

# --------------------------
class SetupPlotFeatures(unittest.TestCase):
    def test_default_features(self):
        nparam = 2
        ns1, ns2, names, figsizeinches = utilities.setup_plot_features(nparam = nparam, names = None, figsizeinches = None)
        self.assertEqual(ns1, math.ceil(math.sqrt(nparam)), msg = 'Expect 3')
        self.assertEqual(ns2, round(math.sqrt(nparam)), msg = 'Expect 2')
        for ii in range(nparam):
            self.assertEqual(names[ii], str('$p_{{{}}}$'.format(ii)))
        self.assertEqual(figsizeinches, [5,4], msg = 'Default figure size is [5,4]')

    def test_nondefault_features(self):
        nparam = 2
        ns1, ns2, names, figsizeinches = utilities.setup_plot_features(nparam = nparam, names = ['hi'], figsizeinches = [7,2])
        self.assertEqual(ns1, math.ceil(math.sqrt(nparam)), msg = 'Expect 3')
        self.assertEqual(ns2, round(math.sqrt(nparam)), msg = 'Expect 2')
        self.assertEqual(names[0], 'hi', msg = 'First name is hi')
        for ii in range(1, nparam):
            self.assertEqual(names[ii], str('$p_{{{}}}$'.format(ii)))
        self.assertEqual(figsizeinches, [7,2], msg = 'Default figure size is [7,2]')

# --------------------------
class GenerateDefaultNames(unittest.TestCase):

    def test_size_of_default_names(self):
        nparam = 8
        names = utilities.generate_default_names(nparam = nparam)
        self.assertEqual(len(names), nparam, msg = 'Length of names should match number of parameters')

    def test_value_of_default_names(self):
        names = utilities.generate_default_names(nparam = 3)
        expected_names = ['$p_{0}$','$p_{1}$','$p_{2}$']
        self.assertEqual(names, expected_names,
                         msg = str('Names do not match: Expected - {}, Received - {}'.format(expected_names, names)))

# --------------------------
class ExtendNamesToMatchNparam(unittest.TestCase):

    def test_initially_empty_name_set(self):
        nparam = 3
        names = utilities.extend_names_to_match_nparam(names = None, nparam = nparam)
        expected_names = ['$p_{0}$','$p_{1}$','$p_{2}$']
        self.assertEqual(names, expected_names,
                         msg = str('Names do not match: Expected - {}, Received - {}'.format(expected_names, names)))

    def test_single_entry_name_set(self):
        nparam = 3
        names = ['aa']
        names = utilities.extend_names_to_match_nparam(names = names, nparam = nparam)
        expected_names = ['aa','$p_{1}$','$p_{2}$']
        self.assertEqual(names, expected_names,
                         msg = str('Names do not match: Expected - {}, Received - {}'.format(expected_names, names)))

    def test_double_entry_name_set(self):
        nparam = 3
        names = ['aa', 'zz']
        names = utilities.extend_names_to_match_nparam(names = names, nparam = nparam)
        expected_names = ['aa','zz','$p_{2}$']
        self.assertEqual(names, expected_names,
                         msg = str('Names do not match: Expected - {}, Received - {}'.format(expected_names, names)))

# --------------------------
class MakeXGrid(unittest.TestCase):

    def test_shape_of_output(self):
        x = np.linspace(0, 10, num = 50)
        npts = 20
        xgrid = utilities.make_x_grid(x = x, npts = 20)
        self.assertEqual(xgrid.shape, (npts, 1), msg = str('Expected return dimension of ({}, 1)'.format(npts)))

    def test_default_shape_of_output(self):
        x = np.linspace(0, 10, num = 50)
        xgrid = utilities.make_x_grid(x = x)
        self.assertEqual(xgrid.shape, (100, 1), msg = 'Expected return dimension of (100, 1)')

    def test_shape_of_output_for_dense_x(self):
        x = np.linspace(0, 10, num = 500)
        npts = 20
        xgrid = utilities.make_x_grid(x = x, npts = 20)
        self.assertEqual(xgrid.shape, (npts, 1), msg = 'Expected return dimension of (npts, 1)')

    def test_default_shape_of_output_for_dense_x(self):
        x = np.linspace(0, 10, num = 500)
        xgrid = utilities.make_x_grid(x = x)
        self.assertEqual(xgrid.shape, (100, 1), msg = 'Expected return dimension of (100, 1)')

# --------------------------
class GenerateEllipse(unittest.TestCase):

    def test_does_non_square_matrix_return_error(self):
        cmat = np.zeros([3,2])
        mu = np.zeros([2,1])
        with self.assertRaises(SystemExit):
            utilities.generate_ellipse(mu, cmat)

    def test_does_non_symmetric_matrix_return_error(self):
        cmat = np.array([[3,2],[1,3]])
        mu = np.zeros([2,1])
        with self.assertRaises(SystemExit):
            utilities.generate_ellipse(mu, cmat)

    def test_does_non_positive_definite_matrix_return_error(self):
        cmat = np.zeros([2,2])
        mu = np.zeros([2,1])
        with self.assertRaises(SystemExit):
            utilities.generate_ellipse(mu, cmat)

    def test_does_good_matrix_return_equal_sized_xy_arrays(self):
        cmat = np.eye(2)
        mu = np.zeros([2,1])
        x,y = utilities.generate_ellipse(mu, cmat)
        self.assertEqual(x.shape,y.shape)

    def test_does_good_matrix_return_correct_size_array(self):
        cmat = np.eye(2)
        mu = np.zeros([2,1])
        ndp = 50 # number of points to generate ellipse shape
        x,y = utilities.generate_ellipse(mu, cmat, ndp)
        self.assertEqual(x.size, ndp)
        self.assertEqual(y.size, ndp)

# --------------------------
class GaussianDensityFunction(unittest.TestCase):

    def test_float_return_with_float_input(self):
        self.assertTrue(isinstance(utilities.gaussian_density_function(x = 0.), float),
                             msg = 'Expected float return')

    def test_float_return_with_int_input(self):
        self.assertTrue(isinstance(utilities.gaussian_density_function(x = 0), float),
                             msg = 'Expected float return')

    def test_float_return_with_float_input_at_nondefault_mean(self):
        self.assertTrue(isinstance(utilities.gaussian_density_function(x = 0., mu = 100), float),
                             msg = 'Expected float return')

# --------------------------
class IQrange(unittest.TestCase):

    def test_array_return_with_column_vector_input(self):
        x = np.random.random_sample(size = (100,1))
        q = utilities.iqrange(x = x)
        self.assertTrue(isinstance(q, np.ndarray), msg = 'Expected array return - received {}'.format(type(q)))
        self.assertEqual(q.size, 1, msg = 'Expect single element array')

    def test_array_return_with_row_vector_input(self):
        x = np.random.random_sample(size = (1,100))
        q = utilities.iqrange(x = x)
        self.assertTrue(isinstance(q, np.ndarray), msg = 'Expected array return - received {}'.format(type(q)))
        self.assertEqual(q.size, 1, msg = 'Expect single element array')

# --------------------------
class ScaleBandWidth(unittest.TestCase):

    def test_array_return_with_column_vector_input(self):
        x = np.random.random_sample(size = (100,1))
        s = utilities.scale_bandwidth(x = x)
        self.assertTrue(isinstance(s, np.ndarray), msg = 'Expected array return - received {}'.format(type(s)))
        self.assertEqual(s.size, 1, msg = 'Expect single element array')

    def test_array_return_with_row_vector_input(self):
        x = np.random.random_sample(size = (1,100))
        s = utilities.scale_bandwidth(x = x)
        self.assertTrue(isinstance(s, np.ndarray), msg = 'Expected array return - received {}'.format(type(s)))
        self.assertEqual(s.size, 1, msg = 'Expect single element array')

    @patch('pymcmcstat.plotting.utilities.iqrange', return_value = -1.0)
    def test_array_return_with_iqrange_lt_0(self, mock_iqrange):
        x = np.random.random_sample(size = (1,100))
        s = utilities.scale_bandwidth(x = x)
        self.assertTrue(isinstance(s, np.ndarray), msg = 'Expected array return - received {}'.format(type(s)))
        self.assertEqual(s.size, 1, msg = 'Expect single element array')

    @patch('pymcmcstat.plotting.utilities.iqrange', return_value = 1.0)
    def test_array_return_with_iqrange_gt_0(self, mock_iqrange):
        x = np.random.random_sample(size = (1,100))
        s = utilities.scale_bandwidth(x = x)
        self.assertTrue(isinstance(s, np.ndarray), msg = 'Expected array return - received {}'.format(type(s)))
        self.assertEqual(s.size, 1, msg = 'Expect single element array')

# --------------------------
class AppendToNrowNcolBasedOnShape(unittest.TestCase):
    def test_shape_is_2d(self):
        nrow = []
        ncol = []
        sh = (2,1)
        nrow, ncol = utilities.append_to_nrow_ncol_based_on_shape(sh = sh, nrow = nrow, ncol = ncol)
        self.assertEqual(nrow, [2], msg = 'Expect [2]')
        self.assertEqual(ncol, [1], msg = 'Expect [1]')

    def test_shape_is_1d(self):
        nrow = []
        ncol = []
        sh = (2,)
        nrow, ncol = utilities.append_to_nrow_ncol_based_on_shape(sh = sh, nrow = nrow, ncol = ncol)
        self.assertEqual(nrow, [2], msg = 'Expect [2]')
        self.assertEqual(ncol, [1], msg = 'Expect [1]')

# --------------------------
class ConvertFlagToBoolean(unittest.TestCase):
    def test_boolean_conversion(self):
        self.assertTrue(utilities.convert_flag_to_boolean(flag = 'on'), msg = 'on -> True')
        self.assertFalse(utilities.convert_flag_to_boolean(flag = 'off'), msg = 'off -> False')

# --------------------------
class SetLocalParameters(unittest.TestCase):
    def test_set_local_parameters(self):
        slp = utilities.set_local_parameters
        self.assertTrue(np.array_equal(slp(ii = 0, local = np.array([0, 0])), np.array([True, True])), msg = 'Expect Array [True, True]')
        self.assertTrue(np.array_equal(slp(ii = 0, local = np.array([0, 1])), np.array([True, False])), msg = 'Expect Array [True, False]')
        self.assertTrue(np.array_equal(slp(ii = 0, local = np.array([1, 0])), np.array([False, True])), msg = 'Expect Array [False, True]')

        self.assertTrue(np.array_equal(slp(ii = 1, local = np.array([0, 1])), np.array([True, True])), msg = 'Expect Array [True, True]')
        self.assertTrue(np.array_equal(slp(ii = 1, local = np.array([1, 0])), np.array([True, True])), msg = 'Expect Array [True, True]')

        self.assertTrue(np.array_equal(slp(ii = 1, local = np.array([2, 2])), np.array([False, False])), msg = 'Expect Array [False, False]')
        self.assertTrue(np.array_equal(slp(ii = 2, local = np.array([1, 2])), np.array([False, True])), msg = 'Expect Array [False, True]')

# --------------------------------------------
class Empirical_Quantiles_Test(unittest.TestCase):

    def test_does_default_empirical_quantiles_return_3_element_array(self):
        test_out = utilities.empirical_quantiles(np.random.rand(10,1))
        self.assertEqual(test_out.shape, (3,1), msg = 'Default output shape is (3,1)')

    def test_does_non_default_empirical_quantiles_return_2_element_array(self):
        test_out = utilities.empirical_quantiles(np.random.rand(10,1), p = np.array([0.2, 0.5]))
        self.assertEqual(test_out.shape, (2,1), msg = 'Non-default output shape should be (2,1)')

    def test_empirical_quantiles_should_not_support_list_input(self):
        with self.assertRaises(AttributeError):
            utilities.empirical_quantiles([-1,0,1])

    def test_empirical_quantiles_vector(self):
        out = utilities.empirical_quantiles(np.linspace(10,20, num = 10).reshape(10,1), p = np.array([0.22, 0.57345]))
        exact = np.array([[12.2], [15.7345]])
        comp = np.linalg.norm(out - exact)
        self.assertAlmostEqual(comp, 0)

# --------------------------
class CheckSettings(unittest.TestCase):

    def test_settings_with_user_none(self):
        user_settings = None
        default_settings = dict(a = False, linewidth = 3, marker = dict(markersize = 5, color = 'g'))
        settings = utilities.check_settings(default_settings = default_settings, user_settings = user_settings)
        self.assertEqual(settings, default_settings, msg = str('Expect dictionaries to match: {} neq {}'.format(settings, default_settings)))

    def test_settings_with_subdict(self):
        user_settings = dict(a = True, fontsize = 12)
        default_settings = dict(a = False, linewidth = 3, marker = dict(markersize = 5, color = 'g'))
        settings = utilities.check_settings(default_settings = default_settings, user_settings = user_settings)
        self.assertEqual(settings['a'], user_settings['a'], msg = 'Expect user setting to overwrite')
        self.assertEqual(settings['marker'], default_settings['marker'], msg = 'Expect default to persist')

    def test_settings_with_subdict_user_ow(self):
        user_settings = dict(a = True, fontsize = 12, marker = dict(color = 'b'))
        default_settings = dict(a = False, linewidth = 3, marker = dict(markersize = 5, color = 'g'))
        settings = utilities.check_settings(default_settings = default_settings, user_settings = user_settings)
        self.assertEqual(settings['a'], user_settings['a'], msg = 'Expect user setting to overwrite')
        self.assertEqual(settings['marker']['color'], user_settings['marker']['color'], msg = 'Expect user to overwrite')
        self.assertEqual(settings['marker']['markersize'], default_settings['marker']['markersize'], msg = 'Expect default to persist')

    def test_settings_with_subdict_user_has_new_setting(self):
        user_settings = dict(a = True, fontsize = 12, marker = dict(color = 'b'), linestyle = '--')
        default_settings = dict(a = False, linewidth = 3, marker = dict(markersize = 5, color = 'g'))
        settings = utilities.check_settings(default_settings = default_settings, user_settings = user_settings)
        self.assertEqual(settings['a'], user_settings['a'], msg = 'Expect user setting to overwrite')
        self.assertEqual(settings['marker']['color'], user_settings['marker']['color'], msg = 'Expect user to overwrite')
        self.assertEqual(settings['marker']['markersize'], default_settings['marker']['markersize'], msg = 'Expect default to persist')
        self.assertEqual(settings['linestyle'], user_settings['linestyle'], msg = 'Expect user setting to be added')

# --------------------------
class SetupSubsample(unittest.TestCase):

    def test_max_lt_nsimu(self):
        skip = 1
        maxpoints = 10
        nsimu = 100
        inds = utilities.setup_subsample(skip, maxpoints, nsimu)
        self.assertTrue(isinstance(inds, np.ndarray),
                        msg='Expect numpy array')
        self.assertEqual(inds.shape, (10,), msg='Expect (10,)')

    def test_max_gt_nsimu(self):
        skip = 1
        maxpoints = 1000
        nsimu = 100
        inds = utilities.setup_subsample(skip, maxpoints, nsimu)
        self.assertTrue(isinstance(inds, np.ndarray),
                        msg='Expect numpy array')
        self.assertEqual(inds.shape, (100,), msg='Expect (100,)')
        skip = 3
        maxpoints = 1000
        nsimu = 100
        inds = utilities.setup_subsample(skip, maxpoints, nsimu)
        self.assertTrue(isinstance(inds, np.ndarray),
                        msg='Expect numpy array')
        self.assertEqual(inds.shape, (34,), msg='Expect (34,)')

    def test_max_eq_nsimu(self):
        skip = 1
        maxpoints = 100
        nsimu = 100
        inds = utilities.setup_subsample(skip, maxpoints, nsimu)
        self.assertTrue(isinstance(inds, np.ndarray),
                        msg='Expect numpy array')
        self.assertEqual(inds.shape, (100,), msg='Expect (100,)')


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.1.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # s_stock_short_horizon [<img src="https://www.arpm.co/lab/icons/icon_permalink.png" width=30 height=30 style="display: inline;">](https://www.arpm.co/lab/redirect.php?code=s_stock_short_horizon&codeLang=Python)
# For details, see [here](https://www.arpm.co/lab/redirect.php?permalink=ExerStockShort).

# +
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import datetime as dt
from pandas.plotting import register_matplotlib_converters
register_matplotlib_converters()

from arpym.tools import add_logo
# -

# ## [Input parameters](https://www.arpm.co/lab/redirect.php?permalink=s_stock_short_horizon-parameters)

# day, month and the year of the plotted value
day = 2
month = 9
year = 2015

# ## [Step 0](https://www.arpm.co/lab/redirect.php?permalink=s_stock_short_horizon-implementation-step00): Load data

# loading data from 2015-05-27 to 2015-12-07
path = '../../../databases/global-databases/high-frequency/db_stock_NOK_intraday/'
df_nokia_stock = pd.read_csv(path + 'data.csv',
                             header=0)
# convert column 'date' from string to datetime64
df_nokia_stock['date'] = pd.to_datetime(df_nokia_stock.date, dayfirst=True)

# ## [Step 1](https://www.arpm.co/lab/redirect.php?permalink=s_stock_short_horizon-implementation-step01): Select the data to be plotted

t_first = dt.datetime(year, month, day, 9, 30)  # starting time
t_last = dt.datetime(year, month, day, 16, 0)  # ending time
# select data
v_stock = df_nokia_stock[(df_nokia_stock.date >= t_first) &
                         (df_nokia_stock.date <= t_last)]

# ## Plots

# +
plt.style.use('arpm')
# extract values from dataframe
t = v_stock.date.values
v_t_stock = v_stock.price.values

number_of_xticks = 6
tick_array = np.linspace(0, t.shape[0]-1, number_of_xticks, dtype=int)
myFmt = mdates.DateFormatter('%H:%M:%S')

fig = plt.figure()
plt.plot_date(t, v_t_stock, '-')
plt.gca().xaxis.set_major_formatter(myFmt)
plt.xticks(t[tick_array])
plt.xlabel('Time')
plt.ylabel('Value')
plt.title(f'NOKIA intraday value on {dt.date(year, month, day)}')
add_logo(fig)


import time

import numpy as np

def timed(f):
    start = time.perf_counter()
    for _ in range(100):
        f(10000)
    end = time.perf_counter()
    return end - start

def append_loop(n):
    """Simple loop with append"""
    my_list = []
    for _ in range(n):
        my_list.append(0)

def add_loop(n):
    """Simple loop with +="""
    my_list = []
    for _ in range(n):
        my_list += [0]

def list_comprehension(n):
    """List comprehension"""
    _ = [0 for _ in range(n)]

def integer_multiplication(n):
    """List and integer multiplication"""
    _ = [0] * n

def numpy_array(n):
    _ = np.zeros(n)

fns = [append_loop, add_loop, list_comprehension, integer_multiplication, numpy_array]
print([timed(f) for f in fns])


#!/usr/bin/env python
# coding: utf-8

# # 5 - Medium Level Example - Modeling Carports and Canopies + Sampling accross a Module!
#
# This journal shows how to model a carport or canopy ~ a fixed structure, usually at a high clearance from the ground, with more than one bifacial solar module in the same inclined-plane to create a "shade" for the cars/people below.
#
# We assume that bifacia_radiacne is already installed in yoru computer. This works for bifacial_radiance v.3 release.
#
# These journal outlines 4 useful uses of bifacial_radiance and some tricks:
#
#

#
 Creating the modules in the canopy/carport

#
 Adding extra geometry for the pillars/posts supporting the carport/canopy

#
 Sampling the rear irradiance with more resolution (more sensors)

#
 and hacking the sensor position to obtain an irradiance map of rear-irradiance.

#
 Adding an object to simulate a car with a specific reflectivity.

#

#
# This is what we will create:
# ![Carport Image We will create](../images_wiki/Carport.png)
#
# ### Steps:
#
#

#
  Setup of Variables through Making OCT Axis

#
  Adding the pillars

#
  Analysis of the collector width

#
  Mapping the irradiance througout all the Carport

#
  Adding a "Car"

#

#

#

# ### 1. Setup of Variables through Making OCT Axis
#
# We've done this before a couple times, no new stuff here.
#
# The magic is that, for doing the carport we see in the figure, we are going to do a 4-up configuration of modules (**numpanels**), and we are going to repeat that 4-UP 7 times (**nMods**)

# In[2]:

from bifacial_radiance import *
import numpy as np

# In[4]:

from pathlib import Path
testfolder = str(Path().resolve().parent.parent / 'bifacial_radiance' / 'TEMP')

timestamp = 4020 # Noon, June 17th.
simulationname = 'HotelCarport'

# MakeModule Parameters
moduletype='PrismSolar'
numpanels = 4  # Carport will have 4 modules along the y direction (N-S since we are facing it to the south) .
x = 0.95
y = 1.95
xgap = 0.15 # Leaving 15 centimeters between modules on x direction
ygap = 0.10 # Leaving 10 centimeters between modules on y direction
zgap = 0 # no gap to torquetube.
sensorsy = 10*numpanels  # this will give 70 sensors per module.

# Other default values:

# TorqueTube Parameters
axisofrotationTorqueTube=False
torqueTube = False
cellLevelModule = False

# SceneDict Parameters
gcr = 0.33   # We are only doing 1 row so this doesn't matter
albedo = 0.28  #'concrete'     # ground albedo
clearance_height = 4.3 # m
nMods = 7 # six modules length.
nRows = 1  # only 1 row

azimuth_ang=180 # Facing south
tilt =20 # tilt.

# Now let's run the example

demo = RadianceObj(simulationname,path = testfolder)  # Create a RadianceObj 'object'
demo.setGround(albedo) # input albedo number or material name like 'concrete'.  To see options, run this without any input.
epwfile = demo.getEPW(40.0583,-74.4057) # NJ lat/lon 40.0583° N, 74.4057
metdata = demo.readEPW(epwfile) # read in the EPW weather data from above
demo.gendaylit(4020)  # Use this to simulate only one hour at a time.
# This allows you to "view" the scene on RVU (see instructions below)
# timestam 4020 : Noon, June 17th.
#demo.genCumSky(demo.epwfile) # Use this instead of gendaylit to simulate the whole year

# Making module with all the variables
moduleDict=demo.makeModule(name=moduletype,x=x,y=y,numpanels = numpanels, xgap=xgap, ygap=ygap)
# create a scene with all the variables
sceneDict = {'tilt':tilt,'pitch': round(gcr/moduleDict['sceney'],3),'clearance_height':clearance_height,'azimuth':azimuth_ang, 'module_type':moduletype, 'nMods': nMods, 'nRows': nRows}
scene = demo.makeScene(moduletype=moduletype, sceneDict=sceneDict) #makeScene creates a .rad file with 20 modules per row, 7 rows.
octfile = demo.makeOct(demo.getfilelist())  # makeOct combines all of the ground, sky and object fil|es into a .oct file.

# If you view the Oct file at this point, you should see the array of 7 modules, with 4 modules each along the collector widt.

#

# ### 2. Adding the pillars
#
# We will add 4 pillars at roughly the back and front corners of the structure. Some of the code below is to calculate the positions of where the pillars will be at.
#
# We are calculating the location with some math geometry

# In[5]:

xright= x*4
xleft=  -xright

#centerhubheight = (1.9*3+1.9/2)*np.sin(tilt*np.pi/180)
y2nd = -(y*numpanels/2)*np.cos(tilt*np.pi/180) + (y)*np.cos(tilt*np.pi/180)
y6th=  -(y*numpanels/2)*np.cos(tilt*np.pi/180) + (y*numpanels)*np.cos(tilt*np.pi/180)
z2nd = (y*np.sin(tilt*np.pi/180))+clearance_height
z6th =  (y*numpanels)*np.sin(tilt*np.pi/180)+clearance_height

name='Post1'
text='! genbox black cuteBox 0.5 0.5 {} | xform -t -0.25 -0.25 0 -t {} {} 0'.format(z2nd, xleft, y2nd)
print (text)
customObject = demo.makeCustomObject(name,text)
demo.appendtoScene(radfile=scene.radfiles, customObject=customObject, text="!xform -rz 0")

name='Post2'
text='! genbox black cuteBox 0.5 0.5 {} | xform -t -0.25 -0.25 0 -t {} {} 0'.format(z2nd, xright, y2nd)
customObject = demo.makeCustomObject(name,text)
demo.appendtoScene(scene.radfiles, customObject, '!xform -rz 0')

name='Post3'
text='! genbox black cuteBox 0.5 0.5 {} | xform -t -0.25 -0.25 0 -t {} {} 0'.format(z6th, xright, y6th)
customObject = demo.makeCustomObject(name,text)
demo.appendtoScene(scene.radfiles, customObject, '!xform -rz 0')

name='Post4'
text='! genbox black cuteBox 0.5 0.5 {} | xform -t -0.25 -0.25 0 -t {} {} 0'.format(z6th, xleft, y6th)
customObject = demo.makeCustomObject(name,text)
demo.appendtoScene(scene.radfiles, customObject, '!xform -rz 0')

octfile = demo.makeOct(demo.getfilelist())  # makeOct combines all of the ground, sky and object files into a .oct file.

# ### View the geometry with the posts on :
#
# ## rvu -vf views\front.vp -e .01 -pe 0.4 -vp 3.5 -20 22 HotelCarport.oct
#
# -pe sets the exposure levels, and -vp sets the view point so the carport is centered (at least on my screen. you can play with the values). It should look like this:
#
# ![Carpport with posts](../images_wiki/Carport.png)
#
# The post should be coindient with the corners of the array on the high-end of the carport, and on the low end of the carport they should be between the lowest module and the next one. Cute!
#
#

#
#

# ### 3. Analysis of the collector width
#
# Now let's do some analysis along the slope of the modules. Each result file will contain irradiance for the 4 modules that make up the slope of the carport. You can select which "module" along the row you sample too.
#
# We are also increasign the number of points sampled accross the collector width, with the  variable **sensorsy** passed to **moduleanalysis**

# In[4]:

analysis = AnalysisObj(octfile, demo.name)  # return an analysis object including the scan dimensions for back irradiance
modWanted = 1
rowWanted = 1
frontscan, backscan = analysis.moduleAnalysis(scene, modWanted=modWanted, rowWanted=rowWanted, sensorsy=sensorsy)

analysis.analysis(octfile, simulationname+"Mod1", frontscan, backscan)  # compare the back vs front irradiance
print('Annual bifacial ratio average:  %0.3f' %( sum(analysis.Wm2Back) / sum(analysis.Wm2Front) ) )
print("")

# This is the module analysis and an image of the results file
# ![This is the module analysed.](../images_wiki/Carport_analysis.PNG)
#
# You can repeat the analysis for any other module in the row:
#
#

# Notice we are passing a CUSTOM simulation name so the results are generated in separate csv files.
#

#

# In[5]:

modWanted = 2
rowWanted = 1
frontscan, backscan = analysis.moduleAnalysis(scene, modWanted=modWanted, rowWanted=rowWanted, sensorsy=sensorsy)

analysis.analysis(octfile, simulationname+"Mod2", frontscan, backscan)  # compare the back vs front irradiance
print('Annual bifacial ratio average:  %0.3f' %( sum(analysis.Wm2Back) / sum(analysis.Wm2Front) ) )

modWanted = 3
rowWanted = 1
frontscan, backscan = analysis.moduleAnalysis(scene, modWanted=modWanted, rowWanted=rowWanted, sensorsy=sensorsy)

analysis.analysis(octfile, simulationname+"Mod3", frontscan, backscan)  # compare the back vs front irradiance
print('Annual bifacial ratio average:  %0.3f' %( sum(analysis.Wm2Back) / sum(analysis.Wm2Front) ) )

#

# ### 4. Mapping the irradiance througout all the Carport ("Hack" the sensors)
#
# You can "hack" the sensors starting locations to obtain an irradinace distribution map. This is easier when the modules are facing South, or East/West. Below is an example, you'll have to repeat over all the modules/ all the surface area with as much resolution as you have patience to see edge-effects.

# In[6]:

# HACK Frontscan and Backscan
frontscan['xstart']=-1.2

analysis.analysis(octfile, simulationname+"Mod3_point2", frontscan, backscan)  # compare the back vs front irradiance
print('Annual bifacial ratio average:  %0.3f' %( sum(analysis.Wm2Back) / sum(analysis.Wm2Front) ) )

#

# The printed Annual bifacial ratio does not include cleaning of the sensors for the material. Some of the sensros might fall in the spacing between the modules (ygaps) or in the structures added, torquetubes, etc. For a real bifacial ratio gain, use the load and clean functions in load.py.
#
# (This process might be automated in a future release TBD)
#

#
#

#

# ### 5. Adding a "Car"
#
# Add a surface (just like we added the pillars) with a specific reflectivity to represent a car. If you are doing hourly simulation you can compare how much the irradiance increases with and without the car, and if you keep track of your parking lot comings/goings this could make an interesting toy-problem: how much are your employees contributing to your rear irradiance production?

# In[7]:

name='Car_1'
carpositionx=-2
carpositiony=-1
text='! genbox white_EPDM HondaFit 1.6 4.5 1.5 | xform -t -0.8 -2.25 0 -t {} {} 0'.format(carpositionx, carpositiony)
customObject = demo.makeCustomObject(name,text)
demo.appendtoScene(scene.radfiles, customObject, '!xform -rz 0')

octfile = demo.makeOct(demo.getfilelist())  # makeOct combines all of the ground, sky and object files into a .oct file.

# Viewing with:
# ## rvu -vf views\front.vp -e .01 -pe 0.019 -vp 1.5 -14 15 HotelCarport.oct
#
#
# ![Behold the Honda-fit sized cube](../images_wiki/Carport_with_car.PNG)

# In[ ]:

docs/tutorials/5 - Medium Level Example - Bifacial Carports and Canopies + sampling across a module!.py


"""Fixtures and configuration for doctests."""

import numpy as np
import pytest

import probnum as pn

@pytest.fixture(autouse=True)
def autoimport_packages(doctest_namespace):
    """This fixture 'imports' standard packages automatically in order to avoid
    boilerplate code in doctests"""

    doctest_namespace["pn"] = pn
    doctest_namespace["np"] = np


#!/Users/zhiyang/anaconda3/bin/python3
"""
=============================
Grouped bar chart with labels
=============================

This example shows a how to create a grouped bar chart and how to annotate
bars with labels.

@modified by Zhiyang Ong, January 15, 2020.
+ Added code to make the comparison between three gender groups:
	- Men (default group)
	- Women (another default group)
	- Others (new group)
"""

import matplotlib
import matplotlib.pyplot as plt
import numpy as np

labels = ['G1', 'G2', 'G3', 'G4', 'G5']
men_means = [20, 34, 30, 35, 27]
women_means = [25, 32, 34, 20, 25]
# @modified by Zhiyang Ong, January 15, 2020.
other_means = [18, 31, 40, 27, 20]

x = np.arange(len(labels))  # the label locations
width = 0.3  # the width of the bars

"""
	@modified by Zhiyang Ong, January 15, 2020.
	Center point of each bar per group, G_i, for number of populations:
	+ 2: x - width/2, x + width/2
	+ 3: x - width, x, x + width
	+ 4: x - width
"""
fig, ax = plt.subplots()
"""
	rects1 = ax.bar(x - width/2, men_means, width, label='Men')
	rects2 = ax.bar(x + width/2, women_means, width, label='Women')

	From https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.bar.html,
		default width is 0.8.
"""

"""
rects1 = ax.bar(x - width, men_means, align='edge', width=0.3, label='Men')
rects2 = ax.bar(x, women_means, align='edge', width=0.3, label='Women')
# @modified by Zhiyang Ong, January 15, 2020.
rects3 = ax.bar(x + width, other_means, align='edge', width=0.3, label='Others')
"""

"""
# @modified by Zhiyang Ong, January 15, 2020.
rects1 = ax.bar(x - 0.3, men_means, align='edge', width=0.3, label='Men')
rects2 = ax.bar(x, women_means, align='edge', width=0.3, label='Women')
rects3 = ax.bar(x + 0.3, other_means, align='edge', width=0.3, label='Others')
"""

rects1 = ax.bar(x - width, men_means, width, label='Men')
rects2 = ax.bar(x, women_means, width, label='Women')
rects3 = ax.bar(x + width, other_means, width, label='Others')

# Add some text for labels, title and custom x-axis tick labels, etc.
ax.set_ylabel('Scores')
ax.set_xlabel('Groups and genders')
ax.set_title('Scores by group and gender')
ax.set_xticks(x)
ax.set_xticklabels(labels)
#ax.legend()
# @modified by Zhiyang Ong, January 15, 2020.
#ax.legend(handletextpad=5)
#	No effect. Does not work.

def autolabel(rects):
	"""Attach a text label above each bar in *rects*, displaying its height."""
	for rect in rects:
		height = rect.get_height()
		ax.annotate('{}'.format(height),
			#xy=(rect.get_x() + rect.get_width() / 2, height),
			# @modified by Zhiyang Ong, January 15, 2020.
			xy=(rect.get_x() + rect.get_width() / 3, height),
			xytext=(0, 3),  # 3 points vertical offset
			textcoords="offset points",
			ha='center', va='bottom')

autolabel(rects1)
autolabel(rects2)

fig.tight_layout()

"""
	@modified by Zhiyang Ong, January 15, 2020.
	Save plot in PDF format.
	This save to PDF command has to appear before the .show()
		command;
		else, the PDF fuile would be empty.
	https://github.com/futurestudio/matplotlib-tutorials/tree/master/export_samples

	Reference:
		Norman Peitek, "Matplotlib — Save Plots as File", from
			{\it Future Studio: Future Studio Tutorials},
			{Fellner, P{\"{o}}hls, Peitek GbR}, Magdeburg,
			Saxony-Anhalt, Germany, August 5, 2019.
			Available at: https://futurestud.io/tutorials/matplotlib-save-plots-as-file;
				last accessed on January 15, 2019.
"""
plt.savefig(__file__+".pdf")
plt.show()

#############################################################################
#
# ------------
#
# References
# """"""""""
#
# The use of the following functions, methods and classes is shown
# in this example:

matplotlib.axes.Axes.bar
matplotlib.pyplot.bar
matplotlib.axes.Axes.annotate
matplotlib.pyplot.annotate

"""
	Reference(s):
	+ \cite{MatplotlibDevelopmentTeam2020a}
		- data visualization gallery (with examples, or sample code)

"""


# -*- coding: utf-8 -*-
"""
Spyder Editor

This is a temporary script file.
"""

3+2

a= 'good morning'
b= 'Vietnam'

print(a)
print(b)

c= a + b

print(c)

import os   #sistema operativo
import pandas as pd   #gestionar datframes
import numpy as np     # numeric python vectores
import matplotlib.pyplot as plt  # grafics

#create a dataframe
#Lists are defined i Python whith [], separate by commas

name = ['Marta', 'Thais', 'Enrique', 'Lluna', 'Cristian']
age = [35,34,45,24,22]
gender = ['Female', 'Famale', 'Male', 'Female','Male']

print(name, age, gender)

class2021 = pd.DataFrame({'name' : name, 'age' : age, 'gender' : gender})

#clean up
del (age,gender,name)

del (a)
del (b)
del (c)

class2021.shape   #dimensionalidad de un dataframe

class2021.head(3)
class2021.tail(2)

#QC OK

edad = class2021.age
del(edad)

#Get working directory
os.chdir(r'C:\Users\marta\git\Edem2021MDA\pep')
os.getcwd()

#save dataframe to excel
class2021.to_excel("class2020.xlsx")
class2021.to_csv("class2020.csv")

#Reset all (carefull this is a "magic" funcion then it doesn´t)
reset -f

#Load basiclibraries
import os   #sistema operativo
import pandas as pd   #gestionar datframes
import numpy as np     # numeric python vectores
import matplotlib.pyplot as plt  # grafics

#change working directory
os.chdir(r'C:\Users\marta\git\Edem2021MDA\pep')
os.getcwd()

#to learn path to file:

rentals_2011 = pd.read_csv ('washington_bike_rentals_2011.csv', sep=';', decimal=',')

rentals_2011.shape
rentals_2011.head()
rentals_2011.tail()

#QC OK

#Extra our first plot
import matplotlib.pyplot as plt

#select  the variable to plot

rentals_2011.cnt
np.mean (rentals_2011.cnt)
np.std (rentals_2011.cnt)

rentals_2011.cnt.mean()    #en python
rentals_2011.cnt.describe()

plt.hist(rentals_2011.cnt)    #histograma
rentals_2011.cnt.hist()

#plot
x=rentals_2011.cnt   #si el nombre de la variable no tiene espacios
x=rentals_2011['cnt']   #cuando hay cararcteres especiales, espacios...

plt.hist(x,edgecolor='black', bins=20)
plt.xticks(np.arange(0, 7000, step=1000))
plt.title("Figure 1. Registred rental in Washington")
plt.ylabel('Frecuencia')
plt.xlabel('Number of rentals')
plt.show() #finalizar el gráfico

plt.hist(y,edgecolor='black')
ptl.show()

weather_2011 = pd.read_csv ('weather_washington_2011.csv', sep=';', decimal=',')

weather_2011.shape
weather_2011.dtypes

weather_2011.head()
weather_2011.tail()

#QC OK

del(x)

rental_weather_2011 =pd.merge(weather_2011, rentals_2011, on="day")

rental_weather_2011.shape

rental_weather_2011.to_csv("rental_weather_2011.csv")

del rental_weather_2011['dteday_y']    #borrar columna

#cambiar nombre columna
rental_weather_2011 =  rental_weather_2011.rename(columns={"dteday_x": "dteday"})

del (weather_2011)
del (rentals_2011)
del (class2021)

rentals_weather_2012= pd.read_csv ('rentals_weather_2012.csv', sep=';', decimal=',')

rentals_weather_2012.shape
rentals_weather_2012.head()
rentals_weather_2012.tail()

#QC ok

#añadimos filas del segundo fichero

rentals_weather_11_12 = rental_weather_2011.append(rentals_weather_2012, ignore_index=True)

#reordenar las columnas

#rentals_weather_11_12 = rentals_weather_11_12[rental_weather_2011.columns]

del(rental_weather_2011)
del(rentals_weather_2012)

wbr=rentals_weather_11_12
del(rentals_weather_11_12)

#describing a nominal variable
#Numerically
#frecuencies

mytable = wbr.groupby(['weathersit']).size()  #separo en grupos y le digo que me diga cuentos

print(mytable)

mytable.sum()    #tabla, autosumate.

#PErcentages
n=mytable.sum()

mytable2 = (mytable/n)*100
print(mytable2)

mytable3 = round(mytable2,1)    #redondear
mytable3

#Barchart1
#lets label the cathegories
bar_list = ['Sunny', 'Cloudy', 'Rainy']
plt.bar(bar_list, mytable2, edgecolor='black')
plt.bar(bar_list, mytable2)
plt.ylabel('Porcentage')
plt.xlabel('Weather')
plt.title('Figure1. Porcentaged weather situacions')
plt.text(1.7, 50,'n: 731')

plt.savefig('bar1.eps')  #Editable en Adove Ilustrator
plt.savefig('bar1.jpg')  #imagen tipo foto no ediable
plt.savefig('bar1.svg')  #Formato vectorial
plt.show()    #finalizar el gráfico

reset -f

import os   #sistema operativo
import pandas as pd   #gestionar datframes
import numpy as np     # numeric python vectores
import matplotlib.pyplot as plt  # grafics

os.chdir(r'C:\Users\marta\git\Edem2021MDA\pep')
wbr = pd.read_csv('WBR_11_12_denormalized_temp.csv', sep=';', decimal=',')

wbr.shape
wbr.head()
wbr.tail()

wbr.cnt
np.mean (wbr.cnt)
np.std (wbr.cnt)

wbr.cnt.describe()
x = wbr['cnt']

plt.hist(x, bins=10, edgecolor='black')
plt.xticks(np.arange(0, 10000, step=1000))

plt.title("Figure 1. Daily Bicycle rentals in Washinting")
plt.ylabel('Frecuency')
plt.xlabel('Number of rented bicycles')
plt.show() #finalizar el gráfico

res = wbr.cnt.describe()

res["mean"]
print(round (res[1],1))

m = res[1]
sd = res[2]
n = res[0]

plt.hist(x, bins=10, edgecolor='black')
plt.xticks(np.arange(0, 10000, step=1000))
plt.title('Figure 3. Daily Bicycle rentals in Washington DC' '\n''by Capital bikeshare. 2011 - 2012')
plt.ylabel('Frecuency')
plt.xlabel('Number of rented bicycles’)
props = dict(boxstyle='round', facecolor='white’,lw=0.5)
textstr = '$\mathrm{Mean}=%.1f$\n$\mathrm{S.D.}=%.1f$\n$\mathrm{n}=%.0f$'%(m, sd, n)
plt.text (6500,110, textstr , bbox=props)
plt.axvline(x=m,
linewidth=1,
linestyle= ('solid', color="red", label='Mean')


"""
Methods for performing train / test splits.
"""

# -----------------------------------------------------------------------------
# IMPORTS
# -----------------------------------------------------------------------------

from typing import Iterator, Tuple

import numpy as np

# -----------------------------------------------------------------------------
# CLASS DEFINITIONS
# -----------------------------------------------------------------------------

class AlternatingSplit:
    """
    Alternating split cross-validator.

    Provides train / test indices to split data in train / test sets.

    The split is performed in an "alternating" way:
    Assume that ``n_splits=3``. In this case, the samples / data points
    are labeled: `A B C A B C A B C ...` In the first split, all points
    labeled `A` or `B` constitute the training set, and `C` is the test
    (or hold-out) set. In the second split, all points labeled `A` or
    `C` are used for training and `B` is the test split. In the final
    split, `A` is held out and training is performed on `B` and `C`.

    This splitting scheme is useful for HCI / ADI data, because it means
    that the effective field rotation in all splits is the same (using
    standard $k$-fold splitting would---for $k=2$---cut the field
    rotation in the training data in half).

    .. note::
        The syntax and usage is closely based on similar ``sklearn``
        classes such as, e.g., :class:`sklearn.model_selection.KFold`.
    """

    def __init__(self, n_splits: int) -> None:

        # Sanity check: we cannot have less than 1 split
        assert n_splits >= 1, 'n_splits must be a positive integer!'
        self.n_splits = n_splits

    def split(self, X: np.ndarray) -> Iterator[Tuple[np.ndarray, np.ndarray]]:
        """
        Generate indices to split data into training and test set.

        Args:
            X: A 2D numpy array of shape `(n_samples, n_features)` that
                contains the training data.

        Yields:
            A 2-tuple consisting of

            * ``train_idx``: A 1D numpy array containing the training
              set indices for that split.
            * ``test_idx``: A 1D numpy array containing the testing set
              indices for that split.
        """

        # Get the number of samples (= number of rows in X)
        n_samples = X.shape[0]

        # Initialize the array of indices that gets split into train / test
        indices = np.arange(n_samples)

        # If n_splits = 1, we do not need to split. Instead, we simply return
        # the indices right away such that train_idx == test_idx. (This is for
        # compatibility reasons in cases where we do not really want to split
        # the data into training and test.)
        if self.n_splits == 1:
            yield indices, indices
            return

        # Otherwise, generate indices for alternating splitting scheme
        for i in range(self.n_splits):

            test_idx = indices[i :: self.n_splits]
            train_idx = np.setdiff1d(indices, test_idx, assume_unique=True)

            yield train_idx, test_idx


#!/usr/bin/env python3

import unittest
from unittest.mock import patch

import numpy as np

from tmc import points

from tmc.utils import load, get_stdout

module_name="src.diamond"
diamond = load(module_name, "diamond")

def patch_name(m, d):
    import importlib
    parts=d.split(".")
    try:
        getattr(importlib.import_module(m), parts[-1])
        p=".".join([m, parts[-1]])
    except ModuleNotFoundError:
        raise
    except AttributeError:
        if len(parts) == 1:
            raise
        try:
            getattr(importlib.import_module(m), parts[-2])
            p=".".join([m] + parts[-2:])
        except AttributeError:
            if len(parts) == 2:
                raise
            getattr(importlib.import_module(m), parts[-3])
            p=".".join([m] + parts[-3:])
    return p

@points('p02-13.1')
class Diamond(unittest.TestCase):

    def test_type(self):
        d=diamond(3)
        self.assertEqual(d.dtype, int, msg="Incorrect element type!")

    def test_shape(self):
        for n in range(1,10):
            d=diamond(n)
            correct_shape=(2*n-1,)*2
            self.assertEqual(d.shape, correct_shape,
                             msg="Incorrect shape for call 'diamond(%i)'!" % n)

    def test_content(self):
        for n in range(1,10):
            d=diamond(n)
            if n==1:
                size=1
            else:
                size=4*n-4
            self.assertEqual(np.sum(d), size,
                             msg="Incorrect number of 1s for call 'diamond(%i)'!" % n)

    def test_calls(self):
        with patch(patch_name(module_name, 'np.eye'), wraps=np.eye) as peye:
            with patch(patch_name(module_name, 'np.concatenate'), wraps=np.concatenate) as pconcatenate:
                d=diamond(3)
                peye.assert_called()
                pconcatenate.assert_called()

if __name__ == '__main__':
    unittest.main()


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.3.0
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_cartesian:
#
# Cartesian axis settings
# =======================
#
# This section documents features used for modifying Cartesian *x* and *y*
# axis settings, including axis scales, tick locations, and tick label
# formatting. It also documents a handy "dual axes" feature.

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_locators:
#
# Tick locations
# --------------
#
# `Tick locators\
# <https://matplotlib.org/3.2.1/gallery/ticks_and_spines/tick-locators.html>`__
# are used to automatically select sensible tick locations
# based on the axis data limits. In ProPlot, you can change the tick locator
# using the `~proplot.axes.Axes.format` keyword arguments `xlocator`,
# `ylocator`, `xminorlocator`, and `yminorlocator` (or their aliases,
# `xticks`, `yticks`, `xminorticks`, and `yminorticks`). This is powered by
# the `~proplot.constructor.Locator` constructor function.
#
# These keyword arguments can be used to apply built-in matplotlib
# `~matplotlib.ticker.Locator`\ s by their "registered" names (e.g.
# ``xlocator='log'``), to draw ticks every ``N`` data values with
# `~matplotlib.ticker.MultipleLocator` (e.g. ``xlocator=2``), or to tick the
# specific locations in a list using `~matplotlib.ticker.FixedLocator` (just
# like `~matplotlib.axes.Axes.set_xticks` and
# `~matplotlib.axes.Axes.set_yticks`). See
# `~proplot.axes.CartesianAxes.format` and `~proplot.constructor.Locator` for
# details.
#
# To generate lists of tick locations, we recommend using ProPlot's
# `~proplot.utils.arange` function -- it’s basically an *endpoint-inclusive*
# version of `numpy.arange`, which is usually what you'll want in this
# context.

# %%
import proplot as plot
import numpy as np
state = np.random.RandomState(51423)
plot.rc.update(
    facecolor=plot.scale_luminance('powderblue', 1.15),
    linewidth=1, fontsize=10,
    color='dark blue', suptitlecolor='dark blue',
    titleloc='upper center', titlecolor='dark blue', titleborder=False,
)
fig, axs = plot.subplots(nrows=8, axwidth=5, aspect=(8, 1), share=0)
axs.format(suptitle='Tick locators demo')

# Step size for tick locations
axs[0].format(
    xlim=(0, 200), xminorlocator=10, xlocator=30,
    title='MultipleLocator'
)

# Specific list of locations
axs[1].format(
    xlim=(0, 10), xminorlocator=0.1,
    xlocator=[0, 0.3, 0.8, 1.6, 4.4, 8, 8.8, 10],
    title='FixedLocator',
)

# Ticks at numpy.linspace(xmin, xmax, N)
axs[2].format(
    xlim=(0, 10), xlocator=('linear', 21),
    title='LinearLocator',
)

# Logarithmic locator, used automatically for log scale plots
axs[3].format(
    xlim=(1, 100), xlocator='log', xminorlocator='logminor',
    title='LogLocator',
)

# Maximum number of ticks, but at "nice" locations
axs[4].format(
    xlim=(1, 7), xlocator=('maxn', 11),
    title='MaxNLocator',
)

# Index locator, only draws ticks where data is plotted
axs[5].plot(np.arange(10) - 5, state.rand(10), alpha=0)
axs[5].format(
    xlim=(0, 6), ylim=(0, 1), xlocator='index',
    xformatter=[r'$\alpha$', r'$\beta$', r'$\gamma$', r'$\delta$', r'$\epsilon$'],
    title='IndexLocator',
)
plot.rc.reset()

# Hide all ticks
axs[6].format(
    xlim=(-10, 10), xlocator='null',
    title='NullLocator',
)

# Tick locations that cleanly divide 60 minute/60 second intervals
axs[7].format(
    xlim=(0, 2), xlocator='dms', xformatter='dms',
    title='Degree-Minute-Second Locator (requires cartopy)',
)

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_formatters:
#
# Tick labels
# -----------
#
# `Tick formatters\
# <https://matplotlib.org/3.2.1/gallery/ticks_and_spines/tick-formatters.html>`__
# are used to convert floating point numbers to
# nicely-formatted tick labels. In ProPlot, you can change the tick formatter
# using the `~proplot.axes.Axes.format` keyword arguments `xformatter` and
# `yformatter`  (or their aliases, `xticklabels` and `yticklabels`). This is
# powered by the `~proplot.constructor.Formatter` constructor function.
#
# These keyword arguments can be used to apply built-in matplotlib
# `~matplotlib.ticker.Formatter`\ s by their "registered" names (e.g.
# ``xformatter='log'``), to apply a ``%``-style format directive with
# `~matplotlib.ticker.FormatStrFormatter` (e.g. ``xformatter='%.0f'``), or
# to apply custom tick labels with `~matplotlib.ticker.FixedFormatter` (just
# like `~matplotlib.axes.Axes.set_xticklabels` and
# `~matplotlib.axes.Axes.set_yticklabels`). They can also be used
# to apply one of ProPlot's new tick formatters -- for example,
# ``xformatter='deglat'`` to label ticks as the geographic latitude,
# ``xformatter='pi'`` to label ticks as fractions of :math:`\pi`,
# or ``xformatter='sci'`` to label ticks with scientific notation.
# See `~proplot.axes.CartesianAxes.format` and
# `~proplot.constructor.Formatter` for details.
#
# ProPlot also changes the default tick formatter to
# `~proplot.ticker.AutoFormatter`. This class trims trailing zeros by
# default, can be used to *omit tick labels* outside of some data range, and
# can add arbitrary prefixes and suffixes to each label. See
# `~proplot.ticker.AutoFormatter` for details. To disable the trailing
# zero-trimming feature, set :rcraw:`formatter.zerotrim` to ``False``.

# %%
import proplot as plot
import numpy as np
plot.rc.update(
    linewidth=1.2, fontsize=10, facecolor='gray0', figurefacecolor='gray2',
    color='gray8', gridcolor='gray8', titlecolor='gray8', suptitlecolor='gray8',
    titleloc='upper center', titleborder=False,
)
fig, axs = plot.subplots(nrows=9, axwidth=5, aspect=(8, 1), share=0)

# Scientific notation
axs[0].format(xlim=(0, 1e20), xformatter='sci', title='SciFormatter')

# N significant figures for ticks at specific values
axs[1].format(
    xlim=(0, 20), xlocator=(0.0034, 3.233, 9.2, 15.2344, 7.2343, 19.58),
    xformatter=('sigfig', 2), title='SigFigFormatter',  # 2 significant digits
)

# Fraction formatters
axs[2].format(
    xlim=(0, 3 * np.pi), xlocator=np.pi / 4, xformatter='pi', title='FracFormatter',
)
axs[3].format(
    xlim=(0, 2 * np.e), xlocator=np.e / 2, xticklabels='e', title='FracFormatter',
)

# Geographic formatters
axs[4].format(
    xlim=(-90, 90), xlocator=30, xformatter='deglat', title='Latitude Formatter'
)
axs[5].format(
    xlim=(0, 360), xlocator=60, xformatter='deglon', title='Longitude Formatter'
)

# User input labels
axs[6].format(
    xlim=(-1.01, 1), xlocator=0.5,
    xticklabels=['a', 'b', 'c', 'd', 'e'], title='FixedFormatter',
)

# Custom style labels
axs[7].format(
    xlim=(0, 0.001), xlocator=0.0001, xformatter='%.E', title='FormatStrFormatter',
)
axs[8].format(
    xlim=(0, 100), xtickminor=False, xlocator=20,
    xformatter='{x:.1f}', title='StrMethodFormatter',
)
axs.format(ylocator='null', suptitle='Tick formatters demo')
plot.rc.reset()

# %%
import proplot as plot
plot.rc.linewidth = 2
plot.rc.fontsize = 11
locator = [0, 0.25, 0.5, 0.75, 1]
fig, axs = plot.subplots(ncols=2, nrows=2, axwidth=1.5, share=0)

# Formatter comparison
axs[0].format(
    xformatter='scalar', yformatter='scalar', title='Matplotlib formatter'
)
axs[1].format(yticklabelloc='both', title='ProPlot formatter')
axs[:2].format(xlocator=locator, ylocator=locator)

# Limiting the tick range
axs[2].format(
    title='Omitting tick labels', ticklen=5, xlim=(0, 5), ylim=(0, 5),
    xtickrange=(0, 2), ytickrange=(0, 2), xlocator=1, ylocator=1
)

# Setting the wrap range
axs[3].format(
    title='Wrapping the tick range', ticklen=5, xlim=(0, 7), ylim=(0, 6),
    xwraprange=(0, 5), ywraprange=(0, 3), xlocator=1, ylocator=1
)
axs.format(
    ytickloc='both', yticklabelloc='both',
    titlepad='0.5em', suptitle='Default formatters demo'
)
plot.rc.reset()

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_datetime:
#
# Datetime ticks
# --------------
#
# ProPlot can also be used to customize the tick locations and tick label
# format of "datetime" axes. To draw ticks on some particular time unit, just
# use a unit string (e.g. ``xlocator='month'``). To draw ticks every ``N``
# time units, just use a (unit, N) tuple (e.g. ``xlocator=('day', 5)``). For
# `% style formatting
# <https://docs.python.org/3/library/datetime.html#strftime-strptime-behavior>`__
# of datetime tick labels, just use a string containing ``'%'`` (e.g.
# ``xformatter='%Y-%m-%d'``). See `~proplot.axes.CartesianAxes.format`,
# `~proplot.constructor.Locator`, and `~proplot.constructor.Formatter` for
# details.

# %%
import proplot as plot
import numpy as np
plot.rc.update(
    linewidth=1.2, fontsize=10, ticklenratio=0.7,
    figurefacecolor='w', facecolor='pastel blue',
    titleloc='upper center', titleborder=False,
)
fig, axs = plot.subplots(nrows=5, axwidth=6, aspect=(8, 1), share=0)
axs[:4].format(xrotation=0)  # no rotation for these examples

# Default date locator
# This is enabled if you plot datetime data or set datetime limits
axs[0].format(
    xlim=(np.datetime64('2000-01-01'), np.datetime64('2001-01-02')),
    title='Auto date locator and formatter'
)

# Concise date formatter introduced in matplotlib 3.1
axs[1].format(
    xlim=(np.datetime64('2000-01-01'), np.datetime64('2001-01-01')),
    xformatter='concise', title='Concise date formatter',
)

# Minor ticks every year, major every 10 years
axs[2].format(
    xlim=(np.datetime64('2000-01-01'), np.datetime64('2050-01-01')),
    xlocator=('year', 10), xformatter='\'%y', title='Ticks every N units',
)

# Minor ticks every 10 minutes, major every 2 minutes
axs[3].format(
    xlim=(np.datetime64('2000-01-01T00:00:00'), np.datetime64('2000-01-01T12:00:00')),
    xlocator=('hour', range(0, 24, 2)), xminorlocator=('minute', range(0, 60, 10)),
    xformatter='T%H:%M:%S', title='Ticks at specific intervals',
)

# Month and year labels, with default tick label rotation
axs[4].format(
    xlim=(np.datetime64('2000-01-01'), np.datetime64('2008-01-01')),
    xlocator='year', xminorlocator='month',  # minor ticks every month
    xformatter='%b %Y', title='Ticks with default rotation',
)
axs.format(
    ylocator='null', suptitle='Datetime locators and formatters demo'
)
plot.rc.reset()

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_scales:
#
# Changing the axis scale
# -----------------------
#
# "Axis scales" like ``'linear'`` and ``'log'`` control the *x* and *y* axis
# coordinate system. To change the axis scale, simply pass e.g.
# ``xscale='log'`` or ``yscale='log'`` to `~proplot.axes.Axes.format`. This
# is powered by the `~proplot.constructor.Scale` constructor function.
#
# ProPlot also makes several changes to the axis scale API:
#
# * By default, the `~proplot.ticker.AutoFormatter` formatter is used for all
#   axis scales instead of e.g. `~matplotlib.ticker.LogFormatter` for
#   `~matplotlib.scale.LogScale` scales. This can be changed e.g. by passing
#   ``xformatter='log'`` or ``yformatter='log'`` to
#   `~proplot.axes.CartesianAxes.format`.
# * To make its behavior consistent with `~proplot.constructor.Locator` and
#   `~proplot.constructor.Formatter`, the `~proplot.constructor.Scale`
#   constructor function returns instances of `~matplotlib.scale.ScaleBase`,
#   and `~matplotlib.axes.Axes.set_xscale` and
#   `~matplotlib.axes.Axes.set_yscale` now accept these class instances in
#   addition to string names like ``'log'``.
# * While matplotlib axis scales must be instantiated with an
#   `~matplotlib.axis.Axis` instance (for backward compatibility reasons),
#   ProPlot axis scales can be instantiated without the axis instance (e.g.
#   ``plot.LogScale()`` instead of ``plot.LogScale(ax.xaxis)``).
# * The default `subs` for the ``'symlog'`` axis scale is now ``np.arange(1, 10)``,
#   and the default `linthresh` is now ``1``. Also the ``'log'`` and ``'symlog'``
#   axis scales now accept the keywords `base`, `linthresh`, `linscale`, and
#   `subs` rather than keywords with trailing ``x`` or ``y``.

# %%
import proplot as plot
import numpy as np
N = 200
lw = 3
plot.rc.update({
    'linewidth': 1, 'ticklabelweight': 'bold', 'axeslabelweight': 'bold'
})
fig, axs = plot.subplots(ncols=2, nrows=2, axwidth=1.8, share=0)
axs.format(suptitle='Axis scales demo', ytickminor=True)

# Linear and log scales
axs[0].format(yscale='linear', ylabel='linear scale')
axs[1].format(ylim=(1e-3, 1e3), yscale='log', ylabel='log scale')
axs[:2].plot(np.linspace(0, 1, N), np.linspace(0, 1000, N), lw=lw)

# Symlog scale
ax = axs[2]
ax.format(yscale='symlog', ylabel='symlog scale')
ax.plot(np.linspace(0, 1, N), np.linspace(-1000, 1000, N), lw=lw)

# Logit scale
ax = axs[3]
ax.format(yscale='logit', ylabel='logit scale')
ax.plot(np.linspace(0, 1, N), np.linspace(0.01, 0.99, N), lw=lw)
plot.rc.reset()

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_scales_new:
#
# Special axis scales
# -------------------
#
# ProPlot introduces several new axis scales. The ``'cutoff'`` scale (see
# `~proplot.scale.CutoffScale`) is useful when the statistical distribution
# of your data is very unusual. The ``'sine'`` scale (see
# `~proplot.scale.SineLatitudeScale`) scales the axis with a sine function,
# resulting in an *area weighted* spherical latitude coordinate, and the
# ``'mercator'`` scale (see `~proplot.scale.MercatorLatitudeScale`) scales
# the axis with the Mercator projection latitude coordinate. The
# ``'inverse'`` scale (see `~proplot.scale.InverseScale`) can be useful when
# working with spectral data, especially with
# :ref:`"dual" unit axes <ug_dual>`.

# %%
import proplot as plot
import numpy as np
fig, axs = plot.subplots(width=6, nrows=4, aspect=(5, 1), sharex=False)
ax = axs[0]

# Sample data
x = np.linspace(0, 4 * np.pi, 100)
dy = np.linspace(-1, 1, 5)
y1 = np.sin(x)
y2 = np.cos(x)
state = np.random.RandomState(51423)
data = state.rand(len(dy) - 1, len(x) - 1)

# Loop through various cutoff scale options
titles = ('Zoom out of left', 'Zoom into left', 'Discrete jump', 'Fast jump')
args = (
    (np.pi, 3),  # speed up
    (3 * np.pi, 1 / 3),  # slow down
    (np.pi, np.inf, 3 * np.pi),  # discrete jump
    (np.pi, 5, 3 * np.pi)  # fast jump
)
locators = (
    np.pi / 3,
    np.pi / 3,
    np.pi * np.append(np.linspace(0, 1, 4), np.linspace(3, 4, 4)),
    np.pi * np.append(np.linspace(0, 1, 4), np.linspace(3, 4, 4)),
)
for ax, iargs, title, locator in zip(axs, args, titles, locators):
    ax.pcolormesh(x, dy, data, cmap='grays', cmap_kw={'right': 0.8})
    for y, color in zip((y1, y2), ('coral', 'sky blue')):
        ax.plot(x, y, lw=4, color=color)
    ax.format(
        xscale=('cutoff', *iargs), title=title,
        xlim=(0, 4 * np.pi), ylabel='wave amplitude',
        xformatter='pi', xlocator=locator,
        xtickminor=False, xgrid=True, ygrid=False, suptitle='Cutoff axis scales demo'
    )

# %%
import proplot as plot
import numpy as np
plot.rc.reset()
fig, axs = plot.subplots(nrows=2, ncols=3, axwidth=1.7, share=0, order='F')
axs.format(
    collabels=('Power scales', 'Exponential scales', 'Cartographic scales'),
    suptitle='Additional axis scales demo'
)
x = np.linspace(0, 1, 50)
y = 10 * x
state = np.random.RandomState(51423)
data = state.rand(len(y) - 1, len(x) - 1)

# Power scales
colors = ('coral', 'sky blue')
for ax, power, color in zip(axs[:2], (2, 1 / 4), colors):
    ax.pcolormesh(x, y, data, cmap='grays', cmap_kw={'right': 0.8})
    ax.plot(x, y, lw=4, color=color)
    ax.format(
        ylim=(0.1, 10), yscale=('power', power),
        title=f'$x^{{{power}}}$'
    )

# Exp scales
for ax, a, c, color in zip(axs[2:4], (np.e, 2), (0.5, 2), colors):
    ax.pcolormesh(x, y, data, cmap='grays', cmap_kw={'right': 0.8})
    ax.plot(x, y, lw=4, color=color)
    ax.format(
        ylim=(0.1, 10), yscale=('exp', a, c),
        title=f"${(a, 'e')[a == np.e]}^{{{(c, '')[c == 1]}x}}$"
    )

# Geographic scales
n = 20
x = np.linspace(-180, 180, n)
y1 = np.linspace(-85, 85, n)
y2 = np.linspace(-85, 85, n)
data = state.rand(len(x) - 1, len(y2) - 1)
for ax, scale, color in zip(axs[4:], ('sine', 'mercator'), ('coral', 'sky blue')):
    ax.plot(x, y1, '-', color=color, lw=4)
    ax.pcolormesh(x, y2, data, cmap='grays', cmap_kw={'right': 0.8})
    ax.format(
        title=scale.title() + ' y-axis', yscale=scale, ytickloc='left',
        yformatter='deg', grid=False, ylocator=20,
        xscale='linear', xlim=None, ylim=(-85, 85)
    )

# %% [raw] raw_mimetype="text/restructuredtext"
# .. _ug_dual:
#
# Dual unit axes
# --------------
#
# The `~proplot.axes.CartesianAxes.dualx` and
# `~proplot.axes.CartesianAxes.dualy` methods can be used to draw duplicate
# *x* and *y* axes meant to represent *alternate units* in the same
# coordinate range as the "parent" axis. This feature is powered by the
# `~proplot.scale.FuncScale` class.
#
# `~proplot.axes.CartesianAxes.dualx` and `~proplot.axes.CartesianAxes.dualy`
# accept either (1) a single linear forward function, (2) a pair of arbitrary
# forward and inverse functions, or (3) a scale name or scale class instance.
# In the latter case, the scale's transforms are used for the forward and
# inverse functions, and the scale's default locators and formatters are used
# for the default `~proplot.scale.FuncScale` locators and formatters.
#
# In the below examples, we generate dual axes with each of these three methods. Note
# that the "parent" axis scale is now arbitrary -- in the first example shown below,
# we create a `~proplot.axes.CartesianAxes.dualx` axis for an axis scaled by the
# `symlog scale <https://matplotlib.org/3.1.0/gallery/scales/symlog_demo.html>`__.

# %%
import proplot as plot
plot.rc.update({'grid.alpha': 0.4, 'linewidth': 1, 'grid.linewidth': 1})
c1 = plot.scale_luminance('cerulean', 0.5)
c2 = plot.scale_luminance('red', 0.5)
fig, axs = plot.subplots(
    [[1, 1, 2, 2], [0, 3, 3, 0]],
    share=0, aspect=2.2, axwidth=3
)
axs.format(
    suptitle='Duplicate axes with custom transformations',
    xcolor=c1, gridcolor=c1,
    ylocator=[], yformatter=[]
)

# Meters and kilometers
ax = axs[0]
ax.format(xlim=(0, 5000), xlabel='meters')
ax.dualx(
    lambda x: x * 1e-3,
    label='kilometers', grid=True, color=c2, gridcolor=c2
)

# Kelvin and Celsius
ax = axs[1]
ax.format(xlim=(200, 300), xlabel='temperature (K)')
ax.dualx(
    lambda x: x - 273.15,
    label='temperature (\N{DEGREE SIGN}C)', grid=True, color=c2, gridcolor=c2
)

# With symlog parent
ax = axs[2]
ax.format(xlim=(-100, 100), xscale='symlog', xlabel='MegaJoules')
ax.dualx(
    lambda x: x * 1e6,
    label='Joules', formatter='log', grid=True, color=c2, gridcolor=c2
)
plot.rc.reset()

# %%
import proplot as plot
plot.rc.update({'grid.alpha': 0.4, 'linewidth': 1, 'grid.linewidth': 1})
c1 = plot.scale_luminance('cerulean', 0.5)
c2 = plot.scale_luminance('red', 0.5)
fig, axs = plot.subplots(ncols=2, share=0, aspect=0.4, axwidth=1.8)
axs.format(suptitle='Duplicate axes with special transformations')

# Pressure as the linear scale, height on opposite axis (scale height 7km)
ax = axs[0]
ax.format(
    xformatter='null', ylabel='pressure (hPa)',
    ylim=(1000, 10), xlocator=[], ycolor=c1, gridcolor=c1
)
ax.dualy(
    'height', label='height (km)', ticks=2.5, color=c2, gridcolor=c2, grid=True
)

# Height as the linear scale, pressure on opposite axis (scale height 7km)
ax = axs[1]  # span
ax.format(
    xformatter='null', ylabel='height (km)', ylim=(0, 20), xlocator='null',
    grid=True, gridcolor=c2, ycolor=c2
)
ax.dualy(
    'pressure', label='pressure (hPa)', locator=100, color=c1, gridcolor=c1, grid=True,
)
plot.rc.reset()

# %%
import proplot as plot
import numpy as np
plot.rc.margin = 0
c1 = plot.scale_luminance('cerulean', 0.5)
c2 = plot.scale_luminance('red', 0.5)
fig, ax = plot.subplots(aspect=(3, 1), width=6)

# Sample data
cutoff = 1 / 5
x = np.linspace(0.01, 0.5, 1000)  # in wavenumber days
response = (np.tanh(-((x - cutoff) / 0.03)) + 1) / 2  # response func
ax.axvline(cutoff, lw=2, ls='-', color=c2)
ax.fill_between([cutoff - 0.03, cutoff + 0.03], 0, 1, color=c2, alpha=0.3)
ax.plot(x, response, color=c1, lw=2)

# Add inverse scale to top
ax.format(
    xlabel='wavenumber (days$^{-1}$)', ylabel='response', grid=False,
    title='Imaginary response function',
    suptitle='Duplicate axes with wavenumber and period',
)
ax = ax.dualx(
    'inverse', locator='log', locator_kw={'subs': (1, 2, 5)}, label='period (days)'
)
plot.rc.reset()