#!/usr/bin/env python
# coding: utf-8
# +
from tqdm import tqdm
from subprocess import Popen, PIPE
from io import StringIO
from PIL import Image
import pdb
import os
from pathlib import Path
import math
import numpy as np
import pandas as pd
import scipy as sp
from functools import partial
from sklearn.metrics import cohen_kappa_score

# for image preprocessing
import cv2
#import imutils
from torchvision.transforms import ToPILImage
from shutil import copyfile

import torch
from fastai import *
from fastai.core import *
from fastai.basic_data import *
from fastai.basic_train import *
from fastai.torch_core import *
from fastai.callbacks import CSVLogger

from fastai.vision import *
from fastai.vision.learner import create_head, cnn_config, num_features_model

from IPython.core.debugger import set_trace
# -

# +
# this cell is for fast submitting, refer https://www.kaggle.com/c/instant-gratification/discussion/94379#latest-546086
# if you are sure your code is right, keep `fast_commit` to True, so after running this cell, code commission can be
# done and you can submit the kernel for LB quickly.

# but it is recommended to first set `fast_commit` to false first and the kernel will run with only partial data to
# check if there is bug in the code.
fast_commit = False  # commit just to check run OK, we can use this for debugging. Write unittest suit in this jupyter notebook
fast_commit_with_commit_runing_less_data = True  # otherwise just exit

random_world = True
# for disable random seed setting. Also, we can use all training data (no validation) for
final_submission = False
# final submission
do_lr_find = False
try:
    sub = pd.read_csv(
        '../input/aptos2019-blindness-detection/sample_submission.csv')
except:
    sub = pd.read_csv('../input/sample_submission.csv')

submitting_to_LB = False
use_less_train_data = False

if fast_commit:
    if len(sub) < 2000:  # commit, not submit to leaderboard
        sub.to_csv('submission.csv', index=False)  # so we can always submit
        if fast_commit_with_commit_runing_less_data:
            use_less_train_data = True
        else:
            exit()
    else:  # this is real submit for leader board
        submitting_to_LB = True
else:  # pretending/testing for submitting to LB
    submitting_to_LB = True
    #do_lr_find = True
# -

# +
# %reload_ext autoreload
# %autoreload 2
# #!nvidia-smi
# -

# ## Utils functions
# helpful functions, for debugging mainly

# +
# not used...
# we could use the contour information!!
def binary_search_img(start, end, f, steps_to_stop):
    if start >= end:
        return start
    if steps_to_stop <= 0:
        raise RuntimeError("Image is strange, cannot find right box")
    mid = (start+end) // 2
    res = f(mid)
    if res == 0:
        return mid
    return binary_search_img(mid+1, end, f, steps_to_stop-1) if res < 0 else \
        binary_search_img(start, mid-1, f, steps_to_stop-1)
# -
# ## Seeding, data preparation

# +
# The images are actually quite big. We will resize to a much smaller size.

bs = 64  # smaller batch size is better for training, but may take longer
sz = 224  # transformed to this size
# -
# +
# Making pretrained weights work without needing to find the default filename
if not os.path.exists('/tmp/.cache/torch/checkpoints/'):
    os.makedirs('/tmp/.cache/torch/checkpoints/')
get_ipython().system(
    "cp '../input/resnet50/resnet50.pth' '/tmp/.cache/torch/checkpoints/resnet50-19c8e357.pth'")

os.listdir('../input')
# -

# +
#import torchsnooper
#import pysnooper

def seed_everything(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

def prepare_for_party():
    print('Make sure cudnn is enabled:', torch.backends.cudnn.enabled)
    if not random_world:
        SEED = 999
        seed_everything(SEED)

prepare_for_party()

def check_image_list_dist(subfolder):
    path_data = Path('../input/aptos2019-blindness-detection')

    il = ImageList.from_folder(path_data/subfolder)

    img_size_stats = {}
    target_blk = []
    for img in il:
        s = tuple(img.size)
        stat = img_size_stats.setdefault(s, 0)
        img_size_stats[s] = stat + 1
        if s[0] == 480:
            bc = get_diag_black_cnt(img)
            target_blk.append(bc)
            print(bc)
    return img_size_stats, target_blk

test_img_size_stats = {(1958, 2588): 134,
                       (480, 640): 1403,  # 3/4
                       (1736, 2416): 225,
                       (1050, 1050): 69,
                       (1944, 2896): 11,
                       (1110, 1467): 2,
                       (1944, 2592): 6,
                       (614, 819): 45,  # around 3/4
                       (576, 768): 2,
                       (1536, 2048): 28,
                       (1117, 1476): 2,
                       (1764, 2146): 1}
# check_image_list_dist('train_images')
train_img_size_stats = {(2136, 3216): 410,
                        (1958, 2588): 533,
                        (1536, 2048): 351,
                        (2848, 4288): 52,
                        (1736, 2416): 638,
                        (1050, 1050): 974,
                        (1000, 1504): 92,
                        (614, 819): 287,
                        (1226, 1844): 61,
                        (1424, 2144): 28,
                        (1944, 2896): 34,
                        (480, 640): 42,
                        (2588, 3388): 141,
                        (1117, 1476): 14,
                        (1110, 1467): 2,
                        (358, 474): 2,
                        (1764, 2146): 1}
# way too different !!!!
# so we just take care of the types with > 30 imgs
# for test 480*640 images, the get_diag_black_cnt number is (137.77405559515324, 3.369116330637486)

# to preprocess our data: we do this:
# 1. cv2 find the contour
# 2. find the center
# 3. use 3/4 crop, binarysearch, until the black margin we wanted is got (we can count the black pixels around the diagonal)
# 4. done

def get_diag_black_cnt(img):
    assert img.data.dtype == torch.float32
    img_r = img.data[0]  # only need to look R channle, most information

    thresh = img_r.mean()/10  # in tensor the value is in range 0 ~ 1

    def blk_cnt(t):
        return sum((dig < thresh).sum() for dig in [t.diagonal(), t.diagonal(offset=5), t.diagonal(offset=-5)])
    img_r_flip = img_r.flip((1))

    return blk_cnt(img_r) + blk_cnt(img_r_flip)

def get_tensor_diag_black_cnt(img):
    assert img.dtype == torch.float32

    img_r = img[0]  # only need to look R channle, most information

    thresh = img_r.mean()/10  # in tensor the value is in range 0 ~ 1

    def blk_cnt(t):
        return sum((dig < thresh).sum() for dig in [t.diagonal(), t.diagonal(offset=5), t.diagonal(offset=-5)])
    img_r_flip = img_r.flip((1))

    return blk_cnt(img_r) + blk_cnt(img_r_flip)

def fastai_img_2_cv2(img): return (image2np(img.data)*255).astype(np.uint8)

# return list, need to post process, should be around center
def find_center(image, cnt_thresh=20):
    """ Thanks to https://www.pyimagesearch.com/2016/02/01/opencv-center-of-contour/ """
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    thresh = cv2.threshold(blurred, cnt_thresh, 255, cv2.THRESH_BINARY)[1]

    cnts = cv2.findContours(thresh, cv2.RETR_EXTERNAL,
                            cv2.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)
    # loop over the contours
    cs = []
    clen = np.array([len(c) for c in cnts])
    mid = np.argmax(clen)

    for c in [cnts[mid]]:
        # compute the center of the contour
        if len(c) < 5:
            return None
        M = cv2.moments(c)
        cX = int(M["m10"] / M["m00"])
        cY = int(M["m01"] / M["m00"])
        cs.append((cX, cY, c))
    return cs

def find_DR_center(img, thresh=20):
    cs = find_center(fastai_img_2_cv2(img), cnt_thresh=thresh)
    if cs is None:
        return None, None, None
    h, w = img.size
    for cx, cy, c in cs:
        if w/3 <= cx <= 2*w/3 and h/3 <= cy <= 2*h/3:
            return cx, cy, c  # width, height correspondingly
    return None, None, None

# should put in a class
def find_box_from_center(img, ratio=3/4, bc=137.774, bc_std=3.369, thresh=20):
    """ratio here, 3 height - 4 width (3y, 4x)
    """
    h, w = img.size

    def contour_range(cx, cy, c):
        xs = c[:, 0, 0]
        ys = c[:, 0, 1]
        xs_rel = xs - cx
        ys_rel = ys - cy
        ratio = xs_rel*3 - ys_rel*4  # minimum one
        sid = np.argsort(np.absolute(ratio))
        xid = sid[0]
        xid_d = None
        for i in sid[1:]:
            if abs(i-xid) < 5:  # just use the trick, as they can't be that close
                continue
            else:
                xid_d = i
                break
        if xid_d is None:
            raise RuntimeError(
                "Image Strange, cannot find right contour points")
        c1, c2 = (xs[xid], ys[xid]), (xs[xid_d], ys[xid_d])
        if c1[0]*c1[1] > c2[0]*c2[0]:
            c2, c1 = c1, c2
        return xs.max(), xs.min(), ys.max(), ys.min(), c1, c2

    cx, cy, c = find_DR_center(img, thresh=thresh)
    if cx is None:
        return None
    xmax, xmin, ymax, ymin, cross, cross2 = contour_range(cx, cy, c)

    def find_box_from_cross_points(c1, c2, ratio, expand=28/400):
        # for 480*640, 137.7 ... we just test more times...
        #138/6 = 23, and the r is around 480/2/3*5 = 800/2=400, so it is about 25/400 more
        # so for h,w, both add these
        width_inner = c2[0] - c1[0]
        width_e = width_inner * expand/2
        hight_e = width_e*ratio
        width_e = int(width_e)
        hight_e = int(hight_e)
        c1[0] -= width_e
        c2[0] += width_e
        c1[1] -= hight_e
        c2[1] += hight_e
        return c1, c2

    def clip(h, w, c):
        x = c[0]
        y = c[1]
        x = x if x > 0 else 0
        y = y if y > 0 else 0
        x = x if x < w else w
        y = y if y < h else h
        return x, y
    c1, c2 = find_box_from_cross_points(list(cross), list(cross2), ratio)
    c1 = clip(h, w, c1)
    c2 = clip(h, w, c2)

    bbox = [c1[1], c1[0], c2[1], c2[0]]

    cnt = get_tensor_diag_black_cnt(img.data[:, c1[1]:c2[1], c1[0]:c2[0]])
    print(cnt)
    if cnt == 0:
        print("Image Strange, cannot find black points in diagonals")

    return ImageBBox.create(h, w, [bbox], labels=[0], classes=['ROI']), bbox
    #ch = cy
    #cw = cx
    #h_range = min(ch-ymin, ymax-ch)  # half box height up limit
    #h_start = abs(ch-cross[1])
    #if h_start > h_range:
    #    raise RuntimeError("Image is strange, cannot find right box")
    #w_range = min(cw-xmin, xmax-cw)
    # start from a more likely place, 2/h_range

    #def check_blk_cnt_center_crop(h_add):
    #    w_add = int(h_add/ratio)
    #    cnt = get_tensor_diag_black_cnt(img.data[:,ch-h_add:ch+h_add,cw-w_add:cw+w_add])
    #    ub = (bc+bc_std)/480
    #    lb = (bc-bc_std)/480
    #    r = cnt.float() / (2*h_add)
    #    if r > ub: return 1
    #    if r < lb: return -1
    #    return 0

    # max_rounds = math.floor(math.log2(h_range-h_start))
    # h_proposed = binary_search_img(h_start, h_range, check_blk_cnt_center_crop, max_rounds)
    # binary search in this situation is not stable, as some image width not enough, so the
    # diagonal calculation will be incorrect
    #
    # easier way, just use the cross points, and add a few pixel margin
    #w_ = int(h_proposed * 4/3)
    #w_ = w_ if w_ < w_range else w_range
    # top left bottem rightb
    #return ImageBBox.create(h,w, [[ch-h_proposed, cw-w_, ch+h_proposed, cw+w_]], labels=[0], classes=['ROI']), [ch-h_proposed, cw-w_, ch+h_proposed, cw+w_]
    #return [ch-h_proposed, cw-w_, ch+h_proposed, cw+w_]

# -
# +
errorsones = []  # to put together and need to handle later

def check_crop(subfolder):
    path_data = Path('../input/aptos2019-blindness-detection')
    cwd = Path('.')
    topil = ToPILImage()

    il = ImageList.from_folder(path_data/subfolder)
    to_save_path = cwd/(subfolder+'_cropped')
    to_save_path_resized = cwd/(subfolder+'_cropped'+'_resized')
    to_save_path_resized.mkdir()
    to_save_path.mkdir()

    cotinue_f = False
    for i, img in enumerate(il):
        fname = il.items[i]
        name_before_resize = to_save_path/fname.name
        name_after_resize = to_save_path_resized/fname.name

        if fname.name != '81914ceb4e74.png' and cotinue_f:
            continue
        if cotinue_f:
            cotinue_f = False
            continue

        s = tuple(img.size)
        if s[0] != 480:
            try:
                _, coords = find_box_from_center(img)
                print(coords, name_before_resize)
                pilimg = topil(
                    img.data[:, coords[0]:coords[2], coords[1]:coords[3]])
                pilimg.save(name_before_resize)
                pilimg_r = pilimg.resize((640, 480))  # weight, height
                pilimg_r.save(name_after_resize)
                del pilimg_r
                del pilimg
            except Exception as e:
                print(e)
                print(fname.name, " error happened")
                errorsones.append(fname.name)
            #break
        else:
            copyfile(fname, name_before_resize)
            copyfile(fname, name_after_resize)

errorsones = []  # to put together and need to handle later

def check_crop_namelist(ns):
    subfolder = 'train_images'
    path_data = Path('../input/aptos2019-blindness-detection')/subfolder
    cwd = Path('.')
    topil = ToPILImage()

    to_save_path = cwd/(subfolder+'_cropped')
    to_save_path_resized = cwd/(subfolder+'_cropped'+'_resized')
    to_save_path_resized.mkdir(parents=True, exist_ok=True)
    to_save_path.mkdir(parents=True, exist_ok=True)

    for imgname in tqdm(ns):
        fname = path_data/imgname
        name_before_resize = to_save_path/fname.name
        name_after_resize = to_save_path_resized/fname.name

        img = open_image(fname)
        s = tuple(img.size)
        if s[0] != 480:
            try:
                # use small thresh, maybe helpful -> much better
                _, coords = find_box_from_center(img, thresh=10)
                print(coords, name_before_resize)
                pilimg = topil(
                    img.data[:, coords[0]:coords[2], coords[1]:coords[3]])
                pilimg.save(name_before_resize)
                pilimg_r = pilimg.resize((640, 480))  # weight, height
                pilimg_r.save(name_after_resize)
                del pilimg_r
                del pilimg
            except Exception as e:
                print(e)
                print(fname.name, " error happened")
                errorsones.append(fname20.name)
            #break
        else:
            copyfile(fname, name_before_resize)
            copyfile(fname, name_after_resize)

still_need_to_handle = [
    '033f2b43de6d.png',
    '3a6e9730b298.png',
    'df4913ca3712.png',
    'f002ce614c59.png',
    '4860f7813654.png',
    'ac720570dd0f.png',
    '807135cbc438.png',
    'e1418d28d668.png',
    '42a67337fa8e.png',
    '17d997fe1090.png',
    '64fedbf97473.png',
    'd51b3fe0fa1b.png',
    '6f4719c6bb4b.png',
    'c58971bcebb2.png',
    '5548a7961a3e.png',
    '7214fc7cbe03.png',
    '7269a1d84a57.png',
    '2c77bf969079.png',
    'e821c1b6417a.png',
    'b665041e1633.png',
    '417f408ee8e0.png',
    '6b7cf869622a.png',
    '66d2ca47aa44.png',
    '6cb96a6fb029.png',
    '50d8a8fb7737.png',
    '541db13517e2.png',
    '5b301a6d1ac7.png',
    'f64214bed40e.png',
    '2bb3c492d6d3.png']
#check_crop_namelist(still_need_to_handle)
# -

def prepare_train_dev_df(df, cls_overlap="None_for_fine"):
    df['path'] = df['id_code'].map(
        lambda x: os.path.join(train_dir, '{}.png'.format(x)))
    df = df.drop(columns=['id_code'])
    df = df.sample(frac=1).reset_index(drop=True)  # shuffle dataframe
    #['normal', 'NPDR_1', 'NPDR_2', 'NPDR_3', 'PDR']
    DR_fine_one_hot = pd.get_dummies(df['diagnosis'], prefix='DR')
    df = pd.concat([df, DR_fine_one_hot], axis=1)

    df['val'] = False
    for i in range(5):
        name = f'DR_{i}'
        val_sub_part = (df[name][df[name] == 1]).sample(
            frac=0.2, replace=False)
        df.loc[val_sub_part.index, 'val'] = True

    if cls_overlap == 'None_for_fine':  # to preserve more information
        pass
    else:
        # could use label smooth, or mixup
        df['DR_3'][df['diagnosis'] == 4] = 1
        df['DR_1'][df['DR_2'] == 1] = 1
        df['DR_1'][df['DR_3'] == 1] = 1
        df['DR_2'][df['DR_3'] == 1] = 1

    NPDR_index = (df['diagnosis'] < 4) & (df['diagnosis'] > 0)

    df['diagnosis_coarse'] = 1
    df['diagnosis_coarse'][df['diagnosis'] == 0] = 0  # Normal
    df['diagnosis_coarse'][NPDR_index] = 1            # NPDR
    df['diagnosis_coarse'][df['diagnosis'] == 4] = 2  # PDR

    DR_coarse_one_hot = pd.get_dummies(df['diagnosis_coarse'], prefix='DRC')
    #DR_coarse_one_hot['DRC_1'][DR_coarse_one_hot['DRC_2'] == 1] = 1

    # score_index = NPDR_index | (df['diagnosis']==4)
    # df['NPDR_score'] = np.NaN
    # df['NPDR_score'][score_index] = df['diagnosis'][score_index]
    # todo use whole data, i.e., normal + PDR

    df['NPDR_score'] = df['diagnosis']

    df = pd.concat([df, DR_coarse_one_hot], axis=1)

    len_df = len(df)
    print(f"There are {len_df} images")

    return df

# +
base_image_dir = os.path.join('..', 'input/aptos2019-blindness-detection/')
train_dir = os.path.join(base_image_dir, 'train_images/')
df = pd.read_csv(os.path.join(base_image_dir, 'train.csv'))
if use_less_train_data:
    df = df.sample(frac=0.1).copy()

df = prepare_train_dev_df(df)
#newpath = "../input/dr-cropped/dr_train_images_cropped_resized/train_images_cropped_resized/"

def replace_image_cropped(df, from_folder, to_folder):
    # path, f62b8a076833.png is truncted, so remove it
    df['path'] = df['path'].str.replace(from_folder, to_folder, regex=False)
    return df[~df['path'].str.contains('f62b8a076833')]

df = replace_image_cropped(df,
                           'aptos2019-blindness-detection/train_images',
                           #'dr-cropped/dr_train_images_cropped_resized/train_images_cropped_resized')
                           'dr-cropped/dr_train_images_cropped')

# This is actually very small. The [previous competition](https://kaggle.com/c/diabetic-retinopathy-detection) had ~35k images, which supports the idea that pretraining on that dataset may be quite beneficial.

# The dataset is highly imbalanced, with many samples for level 0, and very little for the rest of the levels.

df['diagnosis'].hist(figsize=(10, 5))  # so might be overfitting!!!

# Let's look at an example image:

im = Image.open(df['path'][0])
width, height = im.size
print("1st image size: ", width, height)
#im
# -

# +

class CategoryWithScoreProcessor(MultiCategoryProcessor):
    "`PreProcessor` that create `classes` from `ds.items` and handle the mapping."

    def __init__(self, ds: ItemList, one_hot: bool = False):
        super(CategoryWithScoreProcessor, self).__init__(ds, one_hot)

    def process_one(self, item):
        if self.one_hot or isinstance(item, EmptyLabel):
            return item
        return item.astype(np.int)

    def generate_classes(self, items):
        "Generate classes from `items` by taking the sorted unique values."
        raise RuntimeError(
            'For DR, we handle differently, so we don\'t use this function.')

class DRCategory(Category):
    """
    DRCategory just save the raw data as index
    """

    def __int__(self): return int(self.data[-1])

class DRCategoryListWithScore(MultiCategoryList):
    """
    Data format: e.g. [0,1,0,  1,0,0, 1], first 3 is for coarse classification, second 3 for fine
    classification, the last for regresstion
    """
    _processor = CategoryWithScoreProcessor

    def __init__(self, items: Iterator, classes: Collection = None, label_delim: str = None, one_hot: bool = False, **kwargs):
        super().__init__(items, classes=classes,
                         label_delim=label_delim, one_hot=one_hot, **kwargs)
        #self.loss_func = BCEWithLogitsFlat()
        #self.one_hot = one_hot
        #self.copy_new += ['one_hot']
        self.loss_func = None
        self.thresholds = [0.5, 1.5, 2.5, 3.5]

    def set_thresholds(self, thresholds):
        self.thresholds = thresholds

    def get(self, i):
        o = self.items[i]
        if o is None:
            return None
        return DRCategory(o, self.classes[o[-1]])

    def analyze_pred(self, pred, thresh: float = 0.):
        # need to analyze which method got better result
        p = (pred >= thresh).float()
        p[-1] = pred[-1]
        return p
        #p = pred.clone().detach()
        #p[:6] = 0

        #coarse_predict = pred[:3].argmax()
        #p[coarse_predict] = 1.
        #if coarse_predict == 1:  # NPDR
        #    fine_predict = pred[3:6].argmax()
        #    #type_pred = fine_predict + 1
        #    # todo use reg value to analyze. but need to know the threshold
        #    p[3+fine_predict] = 1.
        #return p
        #else:
        #type_pred = coarse_predict if coarse_predict == 0 else 4
        #    return p[]

    def reconstruct(self, t):
        # in ItemList you can see this
        # def analyze_pred(self, pred:Tensor):
        #     "Called on `pred` before `reconstruct` for additional preprocessing."
        reg = t[-1]
        type_pred = 0
        for thr in self.thresholds:
            if reg > thr:
                type_pred += 1  # todo, check if classification info are useful or not
            else:
                break

        return DRCategory(t, self.classes[type_pred])
        # coarse_predict = t[:3].argmax()
        # if coarse_predict == 1:  # NPDR
        #     fine_predict = t[3:6].argmax()
        #     type_pred = fine_predict + 1
        # else:
        #     type_pred = coarse_predict if coarse_predict == 0 else 4

        # return DRCategory(t, self.classes[type_pred])
# -
# +

src = (ImageList.from_df(df=df, path='./', cols='path')  # get dataset from dataset
       # Splitting the dataset
       .split_by_idx(df[df['val'] == 1].index.tolist())
       # obtain labels from the level column
       .label_from_df(cols=['DRC_0', 'DRC_1', 'DRC_2', 'DR_1', 'DR_2', 'DR_3', 'NPDR_score'], label_cls=partial(DRCategoryListWithScore, classes=['normal', 'NPDR_1', 'NPDR_2', 'NPDR_3', 'PDR'], one_hot=False))
       )  # LabelList = ImageList + LabelList
twenty_per_size = int(df['val'].sum())
val_bs = twenty_per_size if twenty_per_size < 800 else 512

tfms = get_transforms(do_flip=True, flip_vert=True, max_rotate=180, max_warp=None,
                      max_zoom=0.9, p_affine=0.5, max_lighting=0.1,
                      p_lighting=0.5)
data = (src.transform(tfms, size=sz, resize_method=ResizeMethod.CROP, padding_mode='reflection')  # Data augmentation
        .databunch(bs=bs, val_bs=val_bs, num_workers=2)  # DataBunch
        )

def get_train_stats(databunch):
    """
    get mean and std of the training set
    """
    tdl = databunch.train_dl

    stats = None
    n = 0
    for x, _ in tdl:
        if stats is None:
            stats = x.new_zeros((2, 3))

        b_n = x.shape[0]
        n += b_n
        # need to record E(X^2) and E(X), for x^2
        stats[0] += x.mean(dim=(0, 2, 3))*b_n
        stats[1] += x.std(dim=(0, 2, 3))*b_n
        print(stats)
    stats /= n
    return stats

#stats = ([0.4285, 0.2286, 0.0753], [0.2700, 0.1485, 0.0812])  # before crop
stats = ([0.5744, 0.3018, 0.0930], [0.1676, 0.1002, 0.0805])  # after crop

if stats is None:
    stats = get_train_stats(data)
    torch.save(stats, 'data_stats.pkl')

DR_img_stats = stats
#imagenet_stats is very different
# Normalize just to mean, std of the data, as it is way too different with imagenet
data = data.normalize((DR_img_stats[0], DR_img_stats[1]))
# this operation will add a transform to data pipeline
# -

data.show_batch(rows=3, figsize=(7, 6))

# +
train_dev_ratio = len(data.dl(DatasetType.Train).x) / \
    len(data.dl(DatasetType.Valid).x)
assert 3.9 < train_dev_ratio < 4.1  # make sure it splits correctly
# -

# ## Training (Transfer learning)

# The Kaggle competition used the Cohen's quadratically weighted kappa so I have that here to compare. This is a better metric when dealing with imbalanced datasets like this one, and for measuring inter-rater agreement for categorical classification (the raters being the human-labeled dataset and the neural network predictions). Here is an implementation based on the scikit-learn's implementation, but converted to a pytorch tensor, as that is what fastai uses.

# ### metric

# #### Optimize the Metric

# Optimizing the quadratic kappa metric was an important part of the top solutions in the previous competition. Thankfully, @abhishek has already provided code to do this for us. We will use this to improve the score.

# +
class OptimizedRounder(object):
    def __init__(self):
        self.coef_ = 0

    def _kappa_loss(self, coef, X, y, cls_weight):
        X_p = np.copy(X)
        for i, pred in enumerate(X_p):
            if pred < coef[0]:
                X_p[i] = 0
            elif pred >= coef[0] and pred < coef[1]:
                X_p[i] = 1
            elif pred >= coef[1] and pred < coef[2]:
                X_p[i] = 2
            elif pred >= coef[2] and pred < coef[3]:
                X_p[i] = 3
            else:
                X_p[i] = 4
        sample_weight = y.new_zeros(y.size(), dtype=torch.float)
        reg_predict = y.clone().detach()
        inds = [torch.nonzero(reg_predict == type_id).squeeze(1)
                for type_id in range(5)]
        for type_id, ind in enumerate(inds):
            sample_weight[ind] = cls_weight[type_id]

        ll = cohen_kappa_score(y, X_p, weights='quadratic',
                               sample_weight=sample_weight)
        return -ll

    def fit(self, X, y, cls_weight):
        loss_partial = partial(self._kappa_loss, X=X,
                               y=y, cls_weight=cls_weight)
        initial_coef = [0.5, 1.5, 2.5, 3.5]
        self.coef_ = sp.optimize.minimize(
            loss_partial, initial_coef, method='nelder-mead')
        print(-loss_partial(self.coef_['x']))

    def predict(self, X, coef):
        X_p = np.copy(X)
        for i, pred in enumerate(X_p):
            if pred < coef[0]:
                X_p[i] = 0
            elif coef[0] <= pred < coef[1]:
                X_p[i] = 1
            elif coef[1] <= pred < coef[2]:
                X_p[i] = 2
            elif coef[2] <= pred < coef[3]:
                X_p[i] = 3
            else:
                X_p[i] = 4
        return X_p

    def coefficients(self):
        return self.coef_['x']

# -

# +

# for batched data, still is the batch size, (64,7)
def convert_to_normal_pred(pred, thresholds):
    thresh_for_PDR = 3.

    if len(pred.shape) == 1:
        #if pred[2] > 0 and pred[-1] > thresh_for_PDR:
        #    return 4
        #coarse_predict = pred[:3].argmax()
        #if coarse_predict == 1:  # NPDR, logit wrong..., argmax need for softmax... and this will ignore all the PDR thing
        #    fine_predict = pred[3:6].argmax()
        #    return fine_predict + 1  # todo use reg value to analyze. but need to know the threshold
        #else:
        #    return coarse_predict if coarse_predict == 0 else 4

        reg = pred[-1]
        type_pred = 0
        for thr in thresholds:
            if reg > thr:
                type_pred += 1  # todo, check if classification info are useful or not
            else:
                break
        return type_pred
    else:
        #coarse_predict = pred[:,:3].argmax(dim=1)  # just do this, and cross our finger for the PDR v.s. NPDR will predict properly.

        #predict = coarse_predict.clone().detach()

        #NPDR_inds_subset = torch.nonzero(coarse_predict == 1).squeeze(1)  # for multi label... how do we do?
        #if len(NPDR_inds_subset) > 0:
        #    fine_predict = pred[NPDR_inds_subset,3:6].argmax(dim=1)
        #    fine_predict += 1  # todo use reg value to analyze. but need to know the threshold

        #    predict[NPDR_inds_subset] = fine_predict

        #reg_value = pred[:,-1]
        #PDR_logit = pred[:,2]

        #PDR_subset = torch.nonzero((PDR_logit > 0) * (reg_value > thresh_for_PDR)).squeeze(1)
        #predict[PDR_subset] = 4
        reg_predict = pred[:, -1].clone().detach()
        inds = [torch.nonzero(reg_predict > thr).squeeze(1)
                for thr in thresholds]
        predict = reg_predict.new_zeros(reg_predict.size(), dtype=torch.int64)
        for type, ind in enumerate(inds):
            predict[ind] = type+1

        del reg_predict

        return predict

def quadratic_kappa(y_hat, y, cls_weight=None, y_hat_predicted=False):
    # need to convert our answer format

    if not y_hat_predicted:
        coefficients = [0.5, 1.5, 2.5, 3.5]
        y_hat = convert_to_normal_pred(y_hat, coefficients)

    if len(y.shape) < 2:
        target = y
    else:
        target = y[:, -1].int()

    y = target
    sample_weight = y.new_ones(y.size(), dtype=torch.float)

    if cls_weight is not None:
        reg_predict = y.clone().detach()
        inds = [torch.nonzero(reg_predict == type_id).squeeze(1)
                for type_id in range(5)]
        for type_id, ind in enumerate(inds):
            sample_weight[ind] = cls_weight[type_id]

    return torch.tensor(cohen_kappa_score(y_hat, target, weights='quadratic', sample_weight=sample_weight), device='cuda:0')
# -

# ### loss

# +
# modified based https://github.com/DingKe/pytorch_workplace/blob/master/focalloss/loss.py
class DR_FocalLoss(nn.Module):
    def __init__(self, gamma: float = 0., eps=1e-7, from_logits=True,
                 cls_cnt=None, cls_overlap="None_for_fine"):
        super(DR_FocalLoss, self).__init__()
        self.gamma = gamma
        self.eps = eps

        self.from_logits = from_logits

        if cls_cnt is not None:
            normal_cnt = cls_cnt[0]

            NPDR_1_cnt = cls_cnt[1]
            NPDR_2_cnt = cls_cnt[2]
            NPDR_3_cnt = cls_cnt[3]

            PDR_cnt = cls_cnt[4]
        else:
            normal_cnt = 1805

            NPDR_1_cnt = 370
            NPDR_2_cnt = 999
            NPDR_3_cnt = 193

            PDR_cnt = 295

        s = normal_cnt + NPDR_1_cnt + NPDR_2_cnt + NPDR_3_cnt + PDR_cnt
        fine_s = NPDR_1_cnt + NPDR_2_cnt + NPDR_3_cnt

        def cal_neg_coef(a, b, c, s=None):  # s can be passed if there is overlay in abc
            if s is None:
                s = a + b + c
            return [a/(s-a), b/(s-b), c/(s-c)]

        def cal_ratio(a, b, c, d):  # s can be passed if there is overlay in abc
            return [a/b, a/c, a/d]

        #b_n_r, b_npdr_r, b_pdr_r = cal_ratio(normal_cnt, normal_cnt, fine_s, PDR_cnt)
        #b_npdr1_r_not_added, b_npdr2_r_not_added, b_npdr3_r_not_added = \
        #    cal_ratio(normal_cnt, NPDR_1_cnt, NPDR_2_cnt, NPDR_3_cnt)
        #self.coarse_mag = [b_n_r, b_npdr_r, b_pdr_r]  # only [-1] is used
        #self.fine_mag_not_added = [b_npdr1_r_not_added, b_npdr2_r_not_added, b_npdr3_r_not_added]
        base_cnt = NPDR_2_cnt

        self.coarse_mag = cal_ratio(base_cnt, normal_cnt, fine_s, PDR_cnt)
        self.fine_mag_not_added = \
            cal_ratio(base_cnt, NPDR_1_cnt, NPDR_2_cnt, NPDR_3_cnt)
        self.reg_mag = 2.  # overall loss for regression

        # following for alpha balancer
        if cls_overlap == "None_for_fine":
            NPDR_3_added = NPDR_3_cnt
            NPDR_2_added = NPDR_2_cnt
            NPDR_1_added = NPDR_1_cnt
        else:
            NPDR_3_added = NPDR_3_cnt + PDR_cnt  # PDR as NPDR3 ... not very good, anyway
            NPDR_2_added = NPDR_2_cnt + NPDR_3_added
            NPDR_1_added = NPDR_2_cnt + NPDR_3_added + NPDR_1_cnt

        NPDR_cnt_added_PDR = fine_s  # + PDR_cnt
        #a_normal, a_NPDR, a_PDR = \
        self.coarse_a = \
            cal_neg_coef(normal_cnt, NPDR_cnt_added_PDR, PDR_cnt, s)
        self.fine_a = \
            cal_neg_coef(NPDR_1_added, NPDR_2_added,
                         NPDR_3_added, NPDR_cnt_added_PDR)
        # only used in balancing classification
        self.fine_mag = \
            cal_ratio(base_cnt, NPDR_1_added, NPDR_2_added, NPDR_3_added)

        print('mag coef: ',
              [self.coarse_mag[0]] + self.fine_mag_not_added +
              [self.coarse_mag[-1]],
              "\nalpha balancer here will be multiplied by negtive loss part\n",
              self.coarse_a, self.fine_a, self.fine_mag)

    @staticmethod
    def cal_balance_ratio_2_parts(pos, neg):
        pos_ratio = neg/(pos+neg)
        return pos_ratio, 1-pos_ratio, pos*pos_ratio

    def forward(self, input, target, **kwargs):
        # so we have 6 logits and one tensor for regression
        reduction_param = kwargs.get('reduction', 'mean')
        target = target.float()  # otherwise, int*float thing

        reg_score_pred = input[:, -1]
        if self.from_logits:
            cls_input = torch.sigmoid(input[:, :-1])
        else:
            cls_input = input[:, :-1]

        coarse_cls_input = cls_input[:, :3]
        coarse_target = target[:, :3]

        # todo add loss/constraint for NPDR,PDR classfication? might be helpful, after we analyze our errors!!
        coarse_cls_loss = self.focal_binary_loss(coarse_cls_input, coarse_target, self.gamma, self.coarse_a, self.eps,
                                                 self.coarse_mag, reduction=reduction_param)
        if reduction_param == 'none':  # need to squash for coarse and fine classes
            coarse_cls_loss = coarse_cls_loss.sum(dim=1)

        loss = coarse_cls_loss

        NPDR_inds_subset = torch.nonzero(coarse_target[:, 1] > 0).squeeze(1)
        fine_cls_input = cls_input[NPDR_inds_subset, 3:6]
        fine_target = target[NPDR_inds_subset, 3:6]
        fine_cls_loss = self.focal_binary_loss(fine_cls_input, fine_target, self.gamma, self.fine_a, self.eps, self.fine_mag,
                                               reduction=reduction_param)
        if reduction_param == 'none':  # need to squash for coarse and fine classes
            fine_cls_loss = fine_cls_loss.sum(dim=1)
            loss[NPDR_inds_subset] += fine_cls_loss  # size ...
        else:
            loss += fine_cls_loss  # a value

        # for reg loss, we do class balance too
        reg_loss = self.reg_mag * \
            F.smooth_l1_loss(reg_score_pred, target[:, -1], reduction='none')

        NPDR1_inds_subset = torch.nonzero(target[:, -1] == 1).squeeze(1)
        NPDR2_inds_subset = torch.nonzero(target[:, -1] == 2).squeeze(1)
        NPDR3_inds_subset = torch.nonzero(target[:, -1] == 3).squeeze(1)
        reg_loss[NPDR1_inds_subset] *= self.fine_mag_not_added[0]
        reg_loss[NPDR2_inds_subset] *= self.fine_mag_not_added[1]
        reg_loss[NPDR3_inds_subset] *= self.fine_mag_not_added[2]

        #normal_inds_subset = torch.nonzero(coarse_target[:, 0] > 0).squeeze(1)
        #reg_loss[normal_inds_subset] *= a_normal

        PDR_inds_subset = torch.nonzero(coarse_target[:, 2] > 0).squeeze(1)
        reg_loss[PDR_inds_subset] *= self.coarse_mag[-1]

        if reduction_param != 'none':
            reg_loss = torch.mean(
                reg_loss) if reduction_param == 'mean' else torch.sum(reg_loss)
        loss += reg_loss

        return loss

    @staticmethod
    def focal_binary_loss(input, target, gamma, alpha, eps, mag, reduction='mean'):
        y_pred = input
        y = target
        not_y = 1 - y
        not_y_pred = 1 - y_pred

        y_pred = y_pred.clamp(eps, 1. - eps)
        not_y_pred = not_y_pred.clamp(eps, 1. - eps)

        if eps > gamma > -eps:  # 0
            pos_gamma_balancer = 1.
            neg_gamma_balancer = 1.
        elif eps > gamma - 1 > -eps:
            pos_gamma_balancer = not_y_pred
            neg_gamma_balancer = y_pred
        else:
            pos_gamma_balancer = not_y_pred ** gamma
            neg_gamma_balancer = y_pred ** gamma

        # RuntimeError: expected backend CUDA and dtype Float but got backend CUDA and dtype Long
        loss = -                       pos_gamma_balancer * \
            y * torch.log(y_pred)  # cross entropy
        loss += -y.new_tensor(alpha) * neg_gamma_balancer * \
            not_y * torch.log(not_y_pred)
        loss *= loss.new_tensor(mag)

        #if target.requires_grad:  # we don't care about this, so just ignore

        if reduction != 'none':
            return torch.mean(loss) if reduction == 'mean' else torch.sum(loss)
        #else:
        #    expanded_input, expanded_target = torch.broadcast_tensors(input, target)
        #    ret = torch._C._nn.mse_loss(expanded_input, expanded_target, _Reduction.get_enum(reduction))
        #return ret
        return loss
# -

# ### learner (resnet 50)

# +
# **Training:**
#
# We use transfer learning, where we retrain the last layers of a pretrained neural network. I use the ResNet50 architecture trained on the ImageNet dataset, which has been commonly used for pre-training applications in computer vision. Fastai makes it quite simple to create a model and train:
def create_DR_head(nf: int, nc: int, lin_ftrs: Optional[Collection[int]] = None, ps: Floats = 0.5,
                   concat_pool: bool = True, bn_final: bool = False):
    "Model head that takes `nf` features, runs through `lin_ftrs`, and about `nc` classes."
    lin_ftrs = [nf, 512, nc] if lin_ftrs is None else [nf] + lin_ftrs + [nc]
    ps = listify(ps)
    if len(ps) == 1:
        ps = [ps[0]/2] * (len(lin_ftrs)-2) + ps
    actns = [nn.ReLU(inplace=True)] * (len(lin_ftrs)-2) + [None]
    pool = AdaptiveConcatPool2d() if concat_pool else nn.AdaptiveAvgPool2d(1)
    layers = [pool, Flatten()]
    for ni, no, p, actn in zip(lin_ftrs[:-1], lin_ftrs[1:], ps, actns):
        layers += bn_drop_lin(ni, no, True, p, actn)
    if bn_final:
        layers.append(nn.BatchNorm1d(lin_ftrs[-1], momentum=0.01))
    return nn.Sequential(*layers)

def DR_learner(data: DataBunch, base_arch: Callable, cut: Union[int, Callable] = None, pretrained: bool = True,
               lin_ftrs: Optional[Collection[int]] = None, ps: Floats = 0.5, custom_head: Optional[nn.Module] = None,
               split_on: Optional[SplitFuncOrIdxList] = None, bn_final: bool = False, init=nn.init.kaiming_normal_,
               concat_pool: bool = True, attention=False, **kwargs: Any) -> Learner:
    "Build convnet style learner."
    meta = cnn_config(base_arch)

    "Create custom convnet architecture"
    body = create_body(base_arch, pretrained, cut)
    if custom_head is None:  # quadnet head
        nf = num_features_model(nn.Sequential(
            *body.children())) * (2 if concat_pool else 1)
        if attention:
            pass
        head = create_DR_head(nf, data.c+2, lin_ftrs, ps=ps,
                              concat_pool=concat_pool, bn_final=bn_final)
    else:
        head = custom_head

    model = nn.Sequential(body, head)

    learn = Learner(data, model, **kwargs)
    learn.split(split_on or meta['split'])
    if pretrained:
        learn.freeze()
    if init:
        apply_init(model[1], init)
    return learn
# -

# +
cls_cnt = df['diagnosis'].value_counts().sort_index().values
# changed lr decay 2/0.15 + patience=3, do not use focal loss...
fl_normal = DR_FocalLoss(gamma=0., cls_cnt=cls_cnt)
# set_trace() we convert back to ipynb
# learner = DR_learner(data, vision.models.densenet121, metrics=[quadratic_kappa])
#learn = cnn_learner(data, base_arch=models.resnet50, metrics=[quadratic_kappa])

def get_cls_weight(cls_cnt, multiply=[1, 1, 1, 1, 1]):
    cls_weight = cls_cnt[0]/cls_cnt
    # cls_weight = [1,1,1,1,1]
    assert len(cls_cnt) == len(multiply)
    for i, m in enumerate(multiply):
        cls_weight[i] *= m

    return cls_weight

#cls_weight = get_cls_weight(cls_cnt, [0.5, 1.5, 3, 1, 0.5])
cls_weight = get_cls_weight(cls_cnt)

learn = DR_learner(data, vision.models.resnet50, cut=-1, loss_func=fl_normal,
                   metrics=[partial(quadratic_kappa, cls_weight=cls_weight)],
                   callback_fns=[partial(CSVLogger, append=True)])

# -

# ### train

# +
def train_triangular_lr(learn, inner_step=0, cycle_cnt=None, max_lr=None, loss=None):
    if inner_step == 0:
        if loss is not None:
            learn.loss_func = loss
        learn.freeze()
        learn.lr_find()
        learn.recorder.plot(suggestion=True)
    elif inner_step == 1:
        assert max_lr is not None and cycle_cnt is not None
        learn.freeze()
        learn.fit_one_cycle(cycle_cnt, max_lr=max_lr, wd=0.1)
        learn.recorder.plot_losses()
        learn.recorder.plot_metrics()
        learn.save('dr-stage1')
    elif inner_step == 2:
        learn.unfreeze()
        learn.lr_find()
        learn.recorder.plot(suggestion=True)
    elif inner_step == 3:
        assert max_lr is not None and cycle_cnt is not None
        # Min numerical gradient: 1.91E-06
        #learn.fit_one_cycle(6, max_lr=slice(1e-6, 5e-6/5))
        learn.unfreeze()
        learn.fit_one_cycle(cycle_cnt, max_lr=max_lr, wd=0.1)

        learn.recorder.plot_losses()
        learn.recorder.plot_metrics()
        learn.save('dr-stage2')
# -

# +
savedPath = Path('models/stage-2.pth')
if savedPath.exists():
    learn.load('stage-2')
# -

# +
# !echo '#!/bin/sh\n( touch metric.log; tail -f history.csv | while true; do nc -v 23.105.212.181 60020; sleep 1; done ) & ' > log_tel.sh
# !chmod +x log_tel.sh
# #!nc -h || apt install netcat -y

#cmdstr = 'python -m unittest unitTest.PSKenelTest.test_pytorch_model_dev'
cmdstr = 'sh ./log_tel.sh'
#Popen(cmdstr.split(), stdout=PIPE)

# -
# +
if submitting_to_LB and do_lr_find:
    train_triangular_lr(learn, inner_step=0)
# -

# +

def get_max_lr(learn):
    try:
        rec = learn.recorder
        lrs = rec._split_list(rec.lrs, 10, 5)
        losses = rec._split_list(rec.losses, 10, 5)
        losses = [x.item() for x in losses]

        mg = (np.gradient(np.array(losses))).argmin()
        print(f"Min numerical gradient: {lrs[mg]:.2E}")
        return lrs[mg]
    except:
        print("Failed to compute the gradients, there might not be enough points.")
        return None
# -

# +
stage_1_cycle = 4 if not use_less_train_data else 1

max_lr_stage_1 = None
if submitting_to_LB and do_lr_find:
    max_lr_stage_1 = get_max_lr(learn)

if max_lr_stage_1 is None or max_lr_stage_1 < 5e-4:
    max_lr_stage_1 = 1e-2
train_triangular_lr(learn, inner_step=1, cycle_cnt=stage_1_cycle,
                    max_lr=max_lr_stage_1)  # choose 3.31E-02 as suggested
# -

# +
if submitting_to_LB and do_lr_find:
    train_triangular_lr(learn, inner_step=2)  # choose 3.31E-02 as suggested
# -

# +
stage_2_cycle = 6 if not use_less_train_data else 1

max_lr_stage_2 = None
if submitting_to_LB and do_lr_find:
    max_lr_stage_2 = get_max_lr(learn)

if max_lr_stage_2 is None or max_lr_stage_2 < 1e-6:
    max_lr_stage_2 = 5e-5
last_layer_lr_scale_down = 10.

last_layer_max_lr_stage_2 = max_lr_stage_1 / last_layer_lr_scale_down
train_triangular_lr(learn, inner_step=3, cycle_cnt=stage_2_cycle, max_lr=slice(
    max_lr_stage_2, last_layer_max_lr_stage_2))
# -

# +
# we need to subclass our our interpretor, as existed one use argmax to predict class

class DRClassificationInterpretation(ClassificationInterpretation):
    def __init__(self, learn: Learner, preds: Tensor, y_true: Tensor, losses: Tensor, ds_type: DatasetType = DatasetType.Valid,
                 cls_converter: Callable = None):
        super(DRClassificationInterpretation, self).__init__(
            learn, preds, y_true, losses, ds_type)
        assert cls_converter is not None
        self.pred_class = cls_converter(self.preds)
        self.y_true = self.y_true[:, -1]

    @classmethod
    def from_learner(cls, learn: Learner,  ds_type: DatasetType = DatasetType.Valid, cls_converter: Callable = None):
        "Gets preds, y_true, losses to construct base class from a learner"
        preds_res = learn.get_preds(ds_type=ds_type, with_loss=True)
        return cls(learn, *preds_res, cls_converter=cls_converter)
# -
# +
# Let's evaluate our model:

coefficients = [0.5, 1.5, 2.5, 3.5]
interp = DRClassificationInterpretation.from_learner(
    learn, cls_converter=partial(convert_to_normal_pred, thresholds=coefficients))

# +
losses, idxs = interp.top_losses()
len(data.valid_ds) == len(losses) == len(idxs)
# -

# +
interp.preds[idxs[:20]]
# -

# +
interp.plot_top_losses(k=12, heatmap=False)
print(interp.confusion_matrix())
print(cohen_kappa_score(interp.pred_class, interp.y_true, weights='quadratic'))
# /opt/conda/lib/python3.6/site-packages/fastai/vision/learner.py:147
# /opt/conda/lib/python3.6/site-packages/fastai/vision/learner.py
# #!cat -n /opt/conda/lib/python3.6/site-packages/fastai/data_block.py
# -

# +
interp.plot_confusion_matrix(figsize=(12, 12), dpi=98)
# -

# +
# ## TTA
#
# Test-time augmentation, or TTA, is a commonly-used technique to provide a boost in your score, and is very simple to implement. Fastai already has TTA implemented, but it is not the best for all purposes, so I am redefining the fastai function and using my custom version.

def _tta_only(learn: Learner, ds_type: DatasetType = DatasetType.Valid, num_pred: int = 10) -> Iterator[List[Tensor]]:
    "Computes the outputs for several augmented inputs for TTA"
    dl = learn.dl(ds_type)
    ds = dl.dataset
    old = ds.tfms
    aug_tfms = [o for o in learn.data.train_ds.tfms]
    try:
        pbar = master_bar(range(num_pred))
        for i in pbar:
            ds.tfms = aug_tfms
            yield get_preds(learn.model, dl, pbar=pbar)[0]
    finally:
        ds.tfms = old

Learner.tta_only = _tta_only

def _TTA(learn: Learner, beta: float = 0, ds_type: DatasetType = DatasetType.Valid, num_pred: int = 1,
         with_loss: bool = False) -> Tensors:
    "Applies TTA to predict on `ds_type` dataset."
    preds, y = learn.get_preds(ds_type)
    if beta is None:
        all_preds = list(learn.tta_only(ds_type=ds_type, num_pred=num_pred))
        avg_preds = torch.stack(all_preds).mean(0)
        return preds, avg_preds, y
    else:
        if beta == 1.:
            final_preds = preds
        else:
            avg_preds = torch.stack(all_preds).mean(0)
            all_preds = list(learn.tta_only(
                ds_type=ds_type, num_pred=num_pred))
            final_preds = preds * beta + avg_preds * (1 - beta)
        if with_loss:
            with NoneReduceOnCPU(learn.loss_func) as lf:
                loss = lf(final_preds, y)
            return final_preds, y, loss
        return final_preds, y

Learner.TTA = _TTA
# -

# ## predict submission.csv

# +
valid_preds = (interp.preds, interp.y_true)

optR = OptimizedRounder()
# might overfit ...(but at V15 code, without it, performance is really bad)
optR.fit(valid_preds[0][:, -1], valid_preds[1], cls_weight=cls_weight)

coefficients = optR.coefficients()
coefficients
# -

# +
test_predictions = None
sample_df = pd.read_csv(
    '../input/aptos2019-blindness-detection/sample_submission.csv')

def preds_for_test(learn, sample_df, optR, coefficients):
    learn.data.add_test(
        ImageList.from_df(sample_df, '../input/aptos2019-blindness-detection',
                          folder='test_images', suffix='.png'))

    preds, y = learn.TTA(ds_type=DatasetType.Test, beta=1.)
    test_predictions = optR.predict(preds[:, -1], coefficients)
    return test_predictions
# -
# +

if submitting_to_LB:
    if test_predictions is None:
        test_predictions = preds_for_test(learn, sample_df, optR, coefficients)
    sample_df.diagnosis = test_predictions.astype(int)
    sample_df.to_csv('submission.csv', index=False)

    sample_df.head()
# -

# +
if submitting_to_LB:
    sample_df.diagnosis.hist()
# -

# +
# use coefficients to predict
interp.pred_class = torch.tensor(optR.predict(interp.preds[:, -1],
                                              coefficients).astype(np.int),
                                 dtype=torch.int64)
# -
# +
print(interp.confusion_matrix())
print(quadratic_kappa(interp.pred_class, interp.y_true,
                      cls_weight=cls_weight, y_hat_predicted=True))
# -


r"""
Facade Sets
"""
#*****************************************************************************
#  Copyright (C) 2010 Nicolas M. Thiery <nthiery at users.sf.net>
#
#  Distributed under the terms of the GNU General Public License (GPL)
#                  http://www.gnu.org/licenses/
#******************************************************************************

from sage.misc.cachefunc import cached_method
from sage.categories.category_singleton import Category_singleton
from sage.categories.category import Category

class FacadeSets(Category_singleton):
    r"""
    The category of facade sets

    A *facade set* is a parent ``P`` whose elements actually belong to
    some other parent::

        sage: P = Sets().example(); P
        Set of prime numbers (basic implementation)
        sage: p = Sets().example().an_element(); p
        47
        sage: p in P
        True
        sage: p.parent()
        Integer Ring

    Typical use cases include modeling a subset of an existing
    parent::

        sage: Sets().Facades().example()
        An example of facade set: the monoid of positive integers

    or the union of several parents::

        sage: Sets().Facades().example("union")
        An example of a facade set: the integers completed by +-infinity

    or endowing a parent with more (or less!) structure::

        sage: Posets().example("facade")
        An example of a facade poset: the positive integers ordered by divisibility

    Let us consider one of the examples above in detail: the partially ordered
    set `P` of positive integers w.r.t. divisibility order. There are two
    options for representing its elements:

     1. as plain integers
     2. as integers, modified to be aware that their parent is `P`

    The advantage of 1. is that one needs not to do conversions back and
    forth between `P` and `\ZZ`. The disadvantage is that this
    introduces an ambiguity when writing `2 < 3`::

        sage: 2 < 3
        True

    To raise this ambiguity, one needs to explicitely specify the order
    as in `2 <_P 3`::

        sage: P = Posets().example("facade")
        sage: P.lt(2,3)
        False

    In short `P` being a facade parent is one of the programmatic
    counterpart (with e.g. coercions) of the usual mathematical idiom:
    "for ease of notation, we identify an element of `P` with the
    corresponding integer". Too many identifications lead to
    confusion; the lack thereof leads to heavy, if not obfuscated,
    notations. Finding the right balance is an art, and even though
    there are common guidelines, it is ultimately up to the writer to
    choose which identifications to do. This is no different in code.

    .. seealso::

       ::

        sage: Sets().example("facade")
        Set of prime numbers (facade implementation)
        sage: Sets().example("inherits")
        Set of prime numbers
        sage: Sets().example("wrapper")
        Set of prime numbers (wrapper implementation)

    .. rubric:: Specifications

    A parent which is a facade must either:

    - call :meth:`Parent.__init__` using the ``facade`` parameter to
      specify a parent, or tuple thereof.
    - overload the method :meth:`~FacadeSets.ParentMethods.facade_for`.

    .. note:: the concept of facade parents was originally introduced
       in the computer algebra system MuPAD.

    TESTS:

    Check that multiple categories initialisation works (:trac:`13801`)::

        sage: class A(Parent):
        ...     def __init__(self):
        ...         Parent.__init__(self, category=(FiniteEnumeratedSets(),Monoids()), facade=True)
        sage: a = A()
    """

    @cached_method
    def super_categories(self):
        r"""
        Returns the super categories of ``self``, as per
        :meth:`Category.super_categories`.

        EXAMPLES::

            sage: Sets().Facades().super_categories()
            [Category of sets]
        """
        from sage.categories.sets_cat import Sets
        return [Sets()]

    def example(self, choice='subset'):
        r"""
        Returns an example of facade set, as per
        :meth:`Category.example()
        <sage.categories.category.Category.example>`.

        INPUT:

        - ``choice`` -- 'union' or 'subset' (default: 'subset').

        EXAMPLES::

            sage: Sets().Facades().example()
            An example of facade set: the monoid of positive integers
            sage: Sets().Facades().example(choice='union')
            An example of a facade set: the integers completed by +-infinity
            sage: Sets().Facades().example(choice='subset')
            An example of facade set: the monoid of positive integers
        """
        import sage.categories.examples.facade_sets as examples
        if choice == "union":
            return examples.IntegersCompletion()
        elif choice == 'subset':
            return examples.PositiveIntegerMonoid()
        else:
            raise TypeError, "choice should be 'union' or 'subset'"

    class ParentMethods:

        def _element_constructor_(self, element):
            """
            Coerce ``element`` into ``self``

            INPUT:

            - ``element`` -- any object

            This default implementation returns ``element`` if
            ``self`` is a facade for ``parent(element)`. Otherwise it
            attempts in turn to coerce ``element`` into each parent
            ``self`` is a facade for.

            This implementation is only valid for a facade parent
            which models the full union of the parents it is a facade
            for. Other facade parents should redefine
            :meth:`element_constructor` appropriately.

            EXAMPLES::

                sage: S = Sets().Facades().example("union"); S
                An example of a facade set: the integers completed by +-infinity
                sage: S(1)
                1
                sage: S(1/2)
                Traceback (most recent call last):
                ...
                ValueError: Can't coerce `1/2` in any parent `An example of a facade set: the integers completed by +-infinity` is a facade for
                sage: S(2/1)
                2
                sage: S(2/1).parent()
                Integer Ring
                sage: S(int(1))
                1
                sage: S(int(1)).parent()
                Integer Ring

            Facade parents that model strict subsets should redefine
            :meth:`element_constructor`::

                sage: S = Sets().Facades().example(); S
                An example of facade set: the monoid of positive integers
                sage: S(-1)
                Traceback (most recent call last):
                ...
                ValueError: %s should be positive
            """
            if self.is_parent_of(element):
                return element
            else:
                parents = self.facade_for()
                if parents is True:
                    return NotImplementedError
                for parent in self.facade_for():
                    try:
                        return parent(element)
                    except Exception:
                        pass
            raise ValueError, "Can't coerce `%s` in any parent `%s` is a facade for"%(element, self)

        def facade_for(self):
            """
            Returns the parents this set is a facade for

            This default implementation assumes that ``self`` has
            an attribute ``_facade_for``, typically initialized by
            :meth:`Parent.__init__`. If the attribute is not present, the method
            raises a NotImplementedError.

            EXAMPLES::

                sage: S = Sets().Facades().example(); S
                An example of facade set: the monoid of positive integers
                sage: S.facade_for()
                (Integer Ring,)

            Check that :trac:`13801` is corrected::

                sage: class A(Parent):
                ...     def __init__(self):
                ...         Parent.__init__(self, category=Sets(), facade=True)
                sage: a = A()
                sage: a.facade_for()
                Traceback (most recent call last):
                ...
                NotImplementedError: this parent did not specify which parents it is a facade for
            """
            try:
                return self._facade_for
            except AttributeError:
                raise NotImplementedError("this parent did not specify which parents it is a facade for")

        def is_parent_of(self, element):
            """
            Returns whether ``self`` is the parent of ``element``

            INPUT:

            - ``element`` -- any object

            Since ``self`` is a facade domain, this actually tests
            whether the parent of ``element`` is any of the parent
            ``self`` is a facade for.

            EXAMPLES::

                sage: S = Sets().Facades().example(); S
                An example of facade set: the monoid of positive integers
                sage: S.is_parent_of(1)
                True
                sage: S.is_parent_of(1/2)
                False

            This method differs from :meth:`__contains__` in two
            ways.  First, this does not take into account the fact
            that ``self`` may be a strict subset of the parent(s)
            it is a facade for::

                sage: -1 in S, S.is_parent_of(-1)
                (False, True)

            Furthermore, there is no coercion attempted::

                sage: int(1) in S, S.is_parent_of(int(1))
                (True, False)

            .. warning::

               this implementation does not handle facade parents of facade
               parents. Is this a feature we want generically?
            """
            parents = self.facade_for()
            if parents is True:
                return True
            from sage.structure.element import parent
            return parent(element) in parents

        def __contains__(self, element):
            """
            Membership testing

            Returns whether ``element`` is in one of the parents
            ``self`` is a facade for.

            .. warning:: this default implementation is currently
            overriden by :meth:`Parent.__contains__`.

            EXAMPLES::

                sage: S = Sets().Facades().example("union"); S
                An example of a facade set: the integers completed by +-infinity
                sage: 1 in S, -5 in S, oo in S, -oo in S, int(1) in S, 2/1 in S
                (True, True, True, True, True, True)
                sage: 1/2 in S, "bla" in S
                (False, False)
            """
            return any(element in parent for parent in self.facade_for())

        def _an_element_(self):
            """
            Try to return an element of ``self``, as per
            :meth:`Sets.ParentMethods.an_element`.

            For each parent ``self`` is a facade for, this default
            implementation tries the method ``an_element`` until it finds an
            element in ``self``. If none is found raise a
            ``NotImplementedError``.

            EXAMPLES::

                sage: S = Sets().Facades().example(); S
                An example of facade set: the monoid of positive integers
                sage: S.an_element()
                1
            """
            for parent in self.facade_for():
                x = parent.an_element()
                if x in self:
                    return x
            raise NotImplementedError


import pandas as pd
import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin

class MyDummyFeatures(BaseEstimator, TransformerMixin):
    """Dummification of categorical features

    Returns:
        pd.DataFrame: transformed pandas DataFrame.
    """

    def __init__(self):
        pass

    def fit(self, X, y=None):
        return self

    def transform(self, X):

        assert isinstance(X, pd.DataFrame)

        return pd.get_dummies(X, drop_first=True)

Assignment 1 - Regression feature engineering/custom_transformers/my_dummy_features.py


__author__ = "Lui Sheng Jie"
__email__ = "luishengjie@outlook.com"

""" Sequential BFS implementation.
    Sequential implementation of Algorithm 1 Parallel BFS algorithm: High-level overview [1].

    Reference:
    [1] https://www.researchgate.net/publication/220782745_Scalable_Graph_Exploration_on_Multicore_Processors

"""

import numpy as np
import time
from src.load_graph import get_graph, gen_balanced_tree

def get_adjacent_nodes(G, x):
    idx_lst = []
    adj_list = G[x]
    for idx, val in enumerate(adj_list):
        if val == 1:
            idx_lst.append(idx)
    return idx_lst

def bfs_seq(G, target):
    r = 0
    CQ = []

    # Init all values in P to inf
    P = [np.inf for i in range(G.shape[0])]
    # Set root node
    P[r] = 0

    # Enqueue r
    CQ.append(r)

    while len(CQ) != 0:
        # print(f"CQ: {CQ}")
        NQ = []

        for i in range(len(CQ)):
            # Dequeue CQ
            u = CQ.pop(0)
            # For each v adjacent to u
            for v in get_adjacent_nodes(G, u):
                if v == target:
                    return True
                if P[v] == np.inf:
                    P[v] = u
                    NQ.append(v)
        # Swap CQ and NQ
        tmp = NQ
        NQ = CQ
        CQ = tmp
    return False

def main():
    start_time = time.time()
    G  = gen_balanced_tree(4, 5, directed=True)
    print(G.shape)
    # G = get_graph()
    find_node = bfs_seq(G, target=999999)
    print("--- %s seconds ---" % (time.time() - start_time))
    if find_node:
        print(f"Node Found")
    else:
        print(f"Node not Found")

if __name__=='__main__':
    main()


# ---
# jupyter:
#   jupytext:
#     cell_metadata_filter: collapsed,code_folding
#     formats: ipynb,py
#     notebook_metadata_filter: all
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.4'
#       jupytext_version: 1.2.3
#   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.7.6
#   latex_envs:
#     LaTeX_envs_menu_present: true
#     autoclose: false
#     autocomplete: true
#     bibliofile: biblio.bib
#     cite_by: apalike
#     current_citInitial: 1
#     eqLabelWithNumbers: true
#     eqNumInitial: 1
#     hotkeys:
#       equation: Ctrl-E
#       itemize: Ctrl-I
#     labels_anchors: false
#     latex_user_defs: false
#     report_style_numbering: false
#     user_envs_cfg: false
# ---

# # Alternative Combinations of Parameter Values
#
# Please write the names and email addresses of everyone who worked on this notebook on the line below.
#
# YOUR NAMES HERE
#
# ## Introduction
#
# The notebook "Micro-and-Macro-Implications-of-Very-Impatient-HHs" is an exercise that demonstrates the consequences of changing a key parameter of the [cstwMPC](http://econ.jhu.edu/people/ccarroll/papers/cstwMPC) model, the time preference factor $\beta$.
#
# The [REMARK](https://github.com/econ-ark/REMARK) `SolvingMicroDSOPs` reproduces the last figure in the [SolvingMicroDSOPs](http://econ.jhu.edu/people/ccarroll/SolvingMicroDSOPs) lecture notes, which shows that there are classes of alternate values of $\beta$ and $\rho$ that fit the data almost as well as the exact 'best fit' combination.
#
# Inspired by this comparison, this notebook asks you to examine the consequences for:
#
# * The consumption function
# * The distribution of wealth
#
# Of _joint_ changes in $\beta$ and $\rho$ together.
#
# One way you can do this is to construct a list of alternative values of $\rho$ (say, values that range upward from the default value of $\rho$, in increments of 0.2, all the way to $\rho=5$).  Then for each of these values of $\rho$ you will find the value of $\beta$ that leads the same value for target market resources, $\check{m}$.
#
# As a reminder, $\check{m}$ is defined as the value of $m$ at which the optimal value of ${c}$ is the value such that, at that value of ${c}$, the expected level of ${m}$ next period is the same as its current value:
#
# $\mathbb{E}_{t}[{m}_{t+1}] = {m}_{t}$
#
# Other notes:
# * The cstwMPC model solves and simulates the problems of consumers with 7 different values of $\beta$
#    * You should do your exercise using the middle value of $\beta$ from that exercise:
#       * `DiscFac_mean   = 0.9855583`
# * You are likely to run into the problem, as you experiment with parameter values, that you have asked HARK to solve a model that does not satisfy one of the impatience conditions required for the model to have a solution.  Those conditions are explained intuitively in the [TractableBufferStock](http://econ.jhu.edu/people/ccarroll/public/lecturenotes/consumption/TractableBufferStock/) model.  The versions of the impatience conditions that apply to the $\texttt{IndShockConsumerType}$ model can be found in the paper [BufferStockTheory](http://econ.jhu.edu/people/ccarroll/papers/BufferStockTheory), table 2.
#    * The conditions that need to be satisfied are:
#       * The Growth Impatience Condition (GIC)
#       * The Return Impatience Condition (RIC)
# * Please accumulate the list of solved consumers' problems in a list called `MyTypes`
#    * For compatibility with a further part of the assignment below

# + {"code_folding": []}
# This cell merely imports and sets up some basic functions and packages

# %matplotlib inline
import matplotlib.pyplot as plt
from tqdm import tqdm
import numpy as np
from copy import deepcopy

import HARK # Prevents import error from Demos repo
from HARK.utilities import plotFuncs

# + {"code_folding": [0, 4]}
# Import IndShockConsumerType
from HARK.ConsumptionSaving.ConsIndShockModel import IndShockConsumerType

# Define a dictionary with calibrated parameters
cstwMPC_calibrated_parameters = {
    "CRRA":1.0,                    # Coefficient of relative risk aversion
    "Rfree":1.01/(1.0 - 1.0/160.0), # Survival probability,
    "PermGroFac":[1.000**0.25], # Permanent income growth factor (no perm growth),
    "PermGroFacAgg":1.0,
    "BoroCnstArt":0.0,
    "CubicBool":False,
    "vFuncBool":False,
    "PermShkStd":[(0.01*4/11)**0.5],  # Standard deviation of permanent shocks to income
    "PermShkCount":5,  # Number of points in permanent income shock grid
    "TranShkStd":[(0.01*4)**0.5],  # Standard deviation of transitory shocks to income,
    "TranShkCount":5,  # Number of points in transitory income shock grid
    "UnempPrb":0.07,  # Probability of unemployment while working
    "IncUnemp":0.15,  # Unemployment benefit replacement rate
    "UnempPrbRet":None,
    "IncUnempRet":None,
    "aXtraMin":0.00001,  # Minimum end-of-period assets in grid
    "aXtraMax":40,  # Maximum end-of-period assets in grid
    "aXtraCount":32,  # Number of points in assets grid
    "aXtraExtra":[None],
    "aXtraNestFac":3,  # Number of times to 'exponentially nest' when constructing assets grid
    "LivPrb":[1.0 - 1.0/160.0],  # Survival probability
    "DiscFac":0.97,             # Default intertemporal discount factor; dummy value, will be overwritten
    "cycles":0,
    "T_cycle":1,
    "T_retire":0,
    'T_sim':1200,  # Number of periods to simulate (idiosyncratic shocks model, perpetual youth)
    'T_age': 400,
    'IndL': 10.0/9.0,  # Labor supply per individual (constant),
    'aNrmInitMean':np.log(0.00001),
    'aNrmInitStd':0.0,
    'pLvlInitMean':0.0,
    'pLvlInitStd':0.0,
    'AgentCount':10000,
}
# -

# Construct a list of solved consumers' problems, IndShockConsumerType is just a place holder
MyTypes = [IndShockConsumerType(**cstwMPC_calibrated_parameters)]

# ## Simulating the Distribution of Wealth for Alternative Combinations
#
# You should now have constructed a list of consumer types all of whom have the same _target_ level of market resources $\check{m}$.
#
# But the fact that everyone has the same target ${m}$ does not mean that the _distribution_ of ${m}$ will be the same for all of these consumer types.
#
# In the code block below, fill in the contents of the loop to solve and simulate each agent type for many periods.  To do this, you should invoke the methods $\texttt{solve}$, $\texttt{initializeSim}$, and $\texttt{simulate}$ in that order.  Simulating for 1200 quarters (300 years) will approximate the long run distribution of wealth in the population.

for ThisType in tqdm(MyTypes):
    ThisType.solve()
    ThisType.initializeSim()
    ThisType.simulate()

# Now that you have solved and simulated these consumers, make a plot that shows the relationship between your alternative values of $\rho$ and the mean level of assets

# +
# To help you out, we have given you the command needed to construct a list of the levels of assets for all consumers
aLvl_all = np.concatenate([ThisType.aLvlNow for ThisType in MyTypes])

# You should take the mean of aLvl for each consumer in MyTypes, divide it by the mean across all simulations
# and then plot the ratio of the values of mean(aLvl) for each group against the value of $\rho$
# -

# # Interpret
# Here, you should attempt to give an intiutive explanation of the results you see in the figure you just constructed

# ## The Distribution of Wealth...
#
# Your next exercise is to show how the distribution of wealth differs for the different parameter  values

# +
from HARK.utilities import getLorenzShares, getPercentiles

# Finish filling in this function to calculate the Euclidean distance between the simulated and actual Lorenz curves.
def calcLorenzDistance(SomeTypes):
    '''
    Calculates the Euclidean distance between the simulated and actual (from SCF data) Lorenz curves at the
    20th, 40th, 60th, and 80th percentiles.

    Parameters
    ----------
    SomeTypes : [AgentType]
        List of AgentTypes that have been solved and simulated.  Current levels of individual assets should
        be stored in the attribute aLvlNow.

    Returns
    -------
    lorenz_distance : float
        Euclidean distance (square root of sum of squared differences) between simulated and actual Lorenz curves.
    '''
    # Define empirical Lorenz curve points
    lorenz_SCF = np.array([-0.00183091,  0.0104425 ,  0.0552605 ,  0.1751907 ])

    # Extract asset holdings from all consumer types
    aLvl_sim = np.concatenate([ThisType.aLvlNow for ThisType in MyTypes])

    # Calculate simulated Lorenz curve points
    lorenz_sim = getLorenzShares(aLvl_sim,percentiles=[0.2,0.4,0.6,0.8])

    # Calculate the Euclidean distance between the simulated and actual Lorenz curves
    lorenz_distance = np.sqrt(np.sum((lorenz_SCF - lorenz_sim)**2))

    # Return the Lorenz distance
    return lorenz_distance

# -

# ## ...and the Marginal Propensity to Consume
#
# Now let's look at the aggregate MPC.  In the code block below, write a function that produces text output of the following form:
#
# $\texttt{The 35th percentile of the MPC is 0.15623}$
#
# Your function should take two inputs: a list of types of consumers and an array of percentiles (numbers between 0 and 1). It should return no outputs, merely print to screen one line of text for each requested percentile.  The model is calibrated at a quarterly frequency, but Carroll et al report MPCs at an annual frequency. To convert, use the formula:
#
# $\kappa_{Y} \approx 1.0 - (1.0 - \kappa_{Q})^4$

# +
# Write a function to tell us about the distribution of the MPC in this code block, then test it!
# You will almost surely find it useful to use a for loop in this function.
def describeMPCdstn(SomeTypes,percentiles):
    MPC_sim = np.concatenate([ThisType.MPCnow for ThisType in SomeTypes])
    MPCpercentiles_quarterly = getPercentiles(MPC_sim,percentiles=percentiles)
    MPCpercentiles_annual = 1.0 - (1.0 - MPCpercentiles_quarterly)**4

    for j in range(len(percentiles)):
        print('The ' + str(100*percentiles[j]) + 'th percentile of the MPC is ' + str(MPCpercentiles_annual[j]))

describeMPCdstn(MyTypes,np.linspace(0.05,0.95,19))
# -

# # If You Get Here ...
#
# If you have finished the above exercises quickly and have more time to spend on this assignment, for extra credit you can do the same exercise where, instead of exploring the consequences of alternative values of relative risk aversion $\rho$, you should test the consequences of different values of the growth factor $\Gamma$ that lead to the same $\check{m}$.


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.6.0
#   kernelspec:
#     display_name: deep_ml_curriculum
#     language: python
#     name: deep_ml_curriculum
# ---

# Note download data from https://drive.google.com/drive/folders/1EgDN57LDuvlZAwr5-eHWB5CTJ7K9HpDP
#
# Credit to this repo: https://github.com/LukasMosser/geolink_dataset
#
# ## Data Disclaimer
#
# All the data serving as an input to these notebooks was generously donated by GEOLINK
# and is CC-by-SA 4.0
#
# If you use their data please reference their dataset properly to give them credit for their contribution.

# %reload_ext autoreload
# %autoreload 2

import lasio
import matplotlib.pyplot as plt
# %matplotlib inline
import os
from tqdm.auto import tqdm
import pandas as pd
import geopandas as gpd
import numpy as np
from pathlib import Path
from sklearn import preprocessing
from operator import itemgetter

# # in and our directories

data_locations = Path(
    "../../data/raw/geolink_dataset/GEOLINK North sea wells with Lithology interpretation/GEOLINK_Lithology and wells NORTH SEA"
)
data_locations_wellheads = Path("../../data/raw/geolink_dataset/norge_well_heads")
interim_locations = Path("../../data/processed/geolink_norge_dataset/")
interim_locations2 = Path("../../data/interim/geolink_norge_dataset/")

# # load and save as parquet

df_lithology = pd.read_excel(data_locations / "../Lithology code data.xlsx", header=1)[
    :-1
]
df_lithology["Abbreviation"] = pd.to_numeric(df_lithology["Abbreviation"])
df_lithology.to_parquet(
    interim_locations / "geolink_norge_lithology.parquet", compression="gzip"
)
df_lithology

# +
# TODO rename well heads
df_well_tops = pd.concat(
    [
        pd.read_csv(data_locations_wellheads / "wellbore_exploration_all.csv"),
        pd.read_csv(data_locations_wellheads / "wellbore_development_all.csv"),
        pd.read_csv(data_locations_wellheads / "wellbore_other_all.csv"),
    ]
)
df_well_tops["wlbWellboreName_geolink"] = df_well_tops["wlbWellboreName"].str.replace(
    "/", "_"
)

# add dates
date_cols = ["wlbEntryDate", "wlbCompletionDate"]
for c in date_cols:
    df_well_tops[c] = pd.to_datetime(df_well_tops[c])  # .astype('str')

df_well_tops["wlbNsDecDeg"] = df_well_tops["wlbNsDecDeg"].replace(0, np.nan)
df_well_tops["wlbEwDesDeg"] = df_well_tops["wlbEwDesDeg"].replace(0, np.nan)

a = set(df_well_tops.columns)
df_well_tops = df_well_tops.dropna(axis=1, thresh=0.9 * len(df_well_tops))
b = set(df_well_tops.columns)
print("removed", a - b)

# make into geodataframe
df_well_tops = gpd.GeoDataFrame(
    df_well_tops,
    geometry=gpd.points_from_xy(df_well_tops.wlbEwDesDeg, df_well_tops.wlbNsDecDeg),
)
df_well_tops
# -

# ## Las files

# We can now proceed to import these files as las files and get their dataframes and hopefully put them into a data format that is more suited for ML tasks.

# +
if not (interim_locations2 / "geolink_norge_well_logs_raw.parquet").exists():

    # load las files
    well_dataframes = []
    files = sorted(data_locations.glob("*.las"))
    for f in tqdm(files):
        df = lasio.read(f).df()
        df["Well"] = f.stem
        well_dataframes.append(df)

    df_all = pd.concat(well_dataframes)

    df_all["Well"] = df_all["Well"].astype("category")

    # Name lithology
    litho_dict = df_lithology.set_index("Abbreviation")["Lithology"].to_dict()
    df_all["LITHOLOGY_GEOLINK"] = (
        df_all["LITHOLOGY_GEOLINK"].replace(litho_dict).astype("category")
    )

    # unique index
    df_all = df_all.reset_index()  # .set_index(['Well', 'DEPT'])

    df_all.to_parquet(
        interim_locations2 / "geolink_norge_well_logs_raw.parquet", compression="gzip"
    )

df_all = pd.read_parquet(interim_locations2 / "geolink_norge_well_logs_raw.parquet")
df_all
# -

# ## Clean las files

# +
# Clean.

# must have well head
df_all_clean2 = df_all[
    df_all.Well.apply(lambda s: s in set(df_well_tops["wlbWellboreName_geolink"]))
]

# must have lithology
df_all_clean2 = df_all_clean2.dropna(subset=["LITHOLOGY_GEOLINK"])
print("nans", df_all_clean2.isna().mean().sort_values())
# Keep /cols logs that are present>thresh of the time
df_all_clean1 = df_all_clean2.dropna(axis=1, thresh=0.6 * len(df_all_clean2))
print('kept {:%} cols'.format(len(df_all_clean1.columns) / len(df_all_clean2.columns)))
# print("nans", df_all_clean1.isna().mean().sort_values())

# Drop rows with any Nan's
df_all_clean = df_all_clean1.dropna(axis=0, how='any')
print('kept {:%} rows'.format(len(df_all_clean) / len(df_all_clean2)))
df_all_clean
# -

df_all_clean.dropna().Well.value_counts()

df_all_clean[df_all_clean['LITHOLOGY_GEOLINK']=='Marlstone'].Well.value_counts()

# +
# 15_9-12
# -

from deep_ml_curriculum.visualization.well_log import plot_facies, plot_well
well_name="30_4-1"
logs = df_all_clean[df_all_clean2.Well==well_name]
facies = logs['LITHOLOGY_GEOLINK'].astype('category').values
plot_well(well_name,
          logs,
          facies)

from deep_ml_curriculum.visualization.well_log import plot_facies, plot_well
well_name="30_6-11"
logs = df_all_clean[df_all_clean2.Well==well_name]
facies = logs['LITHOLOGY_GEOLINK'].astype('category').values
plot_well(well_name,
          logs,
          facies)

# +
# Split by well name
# wells_val = [
#     "35_11-1",
#     "35_11-10",
#     "35_11-11",
#     "35_11-12",
#     "35_11-13",
#     "35_11-15 S",
#     "35_11-2",
#     "35_11-5",
#     "35_11-6",
#     "35_11-7",
#     "35_12-1",
# ]

wells_test = [
    "34_10-12",
    "34_10-16 R",
    "34_10-17",
    "34_10-19",
    "34_10-21",
    "34_10-23",
    "34_10-33",
    "34_10-35",
    "34_10-5",
    "34_10-7",
    "34_11-1",
    "34_11-2 S",
    "34_11-3 T2",
]
# -

df_all_clean_test = df_all_clean[df_all_clean.Well.apply(lambda s: s in wells_test)]
df_all_clean_train = df_all_clean[
    df_all_clean.Well.apply(lambda s: (s not in wells_test))
]
# assert len(set(df_all_clean_val.Well).intersection(set(df_all_clean_train))) == 0
assert len(set(df_all_clean_test.Well).intersection(set(df_all_clean_train))) == 0
# assert len(set(df_all_clean_test.Well).intersection(set(df_all_clean_val))) == 0
len(df_all_clean_train), len(df_all_clean_test)

df_all_clean_train.to_parquet(
    interim_locations / "geolink_norge_well_logs_train.parquet", compression="gzip"
)
df_all_clean_test.to_parquet(
    interim_locations / "geolink_norge_well_logs_test.parquet", compression="gzip"
)
# df_all_clean_val.to_parquet(
#     interim_locations / "geolink_norge_well_logs_val.parquet", compression="gzip"
# )

df_all_clean

# # Others

df_picks = pd.read_excel(
    data_locations / "../NPD stratigraphic picks north sea.xlsx", header=0
)
df_picks.to_parquet(
    interim_locations / "geolink_norge_picks.parquet", compression="gzip"
)

df_picks

# ## Well heads part 2

# only wells we use
a = sorted(df_all.Well.unique())
df_well_tops = df_well_tops[
    df_well_tops["wlbWellboreName_geolink"].apply(lambda s: s in a)
]

df_well_tops.to_file(interim_locations / "norge_well_tops.gpkg", driver="GPKG")

# # Example Load

# +
# Test load
df_all_clean2 = pd.read_parquet(
    interim_locations / "geolink_norge_well_logs_train.parquet"
)  # .set_index(['Well', 'DEPT'])

df_well_tops = gpd.read_file(interim_locations / "norge_well_tops.gpkg")
df_well_tops_minimal = df_well_tops[
    [
        "wlbWellboreName_geolink",
        "wlbCompletionYear",
        "wlbKellyBushElevation",
        "wlbCompletionDate",
        "wlbTotalDepth",
        "geometry",
    ]
]
df_well_tops.plot()
# -

# Merge well tops and well logs, a selection
df_all_clean3 = pd.merge(
    left=df_all_clean2.sample(1000),
    right=df_well_tops_minimal,
    left_on="Well",
    right_on="wlbWellboreName_geolink",
    how="left",
).drop(columns="wlbWellboreName_geolink")
df_all_clean3 = df_all_clean3.set_index(['Well', 'DEPT'])
df_all_clean3 = gpd.GeoDataFrame(df_all_clean3, geometry=df_all_clean3['geometry'])
df_all_clean3.plot()
# df_all_clean3

df_picks = pd.read_parquet(interim_locations / "geolink_norge_picks.parquet")
df_picks

df_all_clean = pd.read_parquet(
    interim_locations / "geolink_norge_well_logs_train.parquet"
).set_index(["Well", "DEPT"])
df_all_clean

# # Example plot

df_all_clean = pd.read_parquet(
    interim_locations / "geolink_norge_well_logs_train.parquet"
).set_index(["Well", "DEPT"])
df_all_clean['DEPT'] = df_all_clean.index.get_level_values(1)
df_all_clean

# +
# logs
# -

from deep_ml_curriculum.visualization.well_log import plot_facies, plot_well
well_name="30_4-1"
logs = df_all_clean.xs(well_name)
facies = logs['LITHOLOGY_GEOLINK'].astype('category').values
plot_well(well_name,
          logs,
          facies)

plt.figure(figsize=(1,8))
plot_facies(facies, plt.gca(), colorbar=False)

# # reindex depth and to Xarray
#
# This lets us includes location easily without using much more space

# +
# Load some
df_all_clean1 = pd.read_parquet(
    interim_locations / "geolink_norge_well_logs_test.parquet"
).set_index(['Well', 'DEPT'])
df_all_clean1['Depth'] = df_all_clean1.index.get_level_values(1)
df_all_clean1['split'] = 'test'

# Load some
df_all_clean2 = pd.read_parquet(
    interim_locations / "geolink_norge_well_logs_train.parquet"
).set_index(['Well', 'DEPT'])
df_all_clean2['Depth'] = df_all_clean2.index.get_level_values(1)
df_all_clean2['split'] = 'train'

# # Load some
# df_all_clean3 = pd.read_parquet(
#     interim_locations / "geolink_norge_well_logs_val.parquet"
# ).set_index(['Well', 'DEPT'])
# df_all_clean3['Depth'] = df_all_clean3.index.get_level_values(1)
# df_all_clean3['split'] = 'val'

df_all = pd.concat([df_all_clean1, df_all_clean2])
df_all
# -

df_well_tops = gpd.read_file(interim_locations / "norge_well_tops.gpkg")
df_well_tops_minimal = df_well_tops[
    [
        "wlbWellboreName_geolink",
        "wlbCompletionYear",
        "wlbKellyBushElevation",
        "wlbCompletionDate",
        "wlbTotalDepth",
        "geometry",
    ]
].copy()
df_well_tops_minimal['xc'] = df_well_tops_minimal.geometry.x
df_well_tops_minimal['yc'] = df_well_tops_minimal.geometry.y
df_well_tops_minimal

nidx = np.arange(400, 5500, 0.15)

# +
def reindex(x):
    """Reindex each well to 15cm"""
    if len(x)==0: return None
    x = x.reset_index().set_index('DEPT')
    x = x.reindex(nidx, method='nearest', limit=1).drop(columns=['Well']).sort_index()
    return x
#     return x.reset_index().set_index(['Well', 'DEPT'])

df_all3 = df_all.groupby(level=0).apply(reindex).dropna()
df_all3
# -

import xarray as xr
xr_all_clean2 = df_all3.to_xarray()
xr_all_clean2

xr_wells = df_well_tops_minimal.rename(columns={'wlbWellboreName_geolink':'Well'}).set_index('Well').to_xarray()
xr_wells

xr_all = xr.merge(
    [xr_all_clean2, xr_wells],
    join='left')

xr_all2 = xr_all.sortby(['Well', 'DEPT'])
xr_all2

well_name="30_4-1"
logs = xr_all2.sel(Well=well_name).to_dataframe().dropna()
logs['DEPT'] = logs['Depth']
facies = logs['LITHOLOGY_GEOLINK'].astype('category').values
plot_well(well_name, logs, facies)
logs

from deep_ml_curriculum.visualization.well_log import plot_facies, plot_well
well_name="30_4-1"
logs = df_all_clean.xs(well_name)
facies = logs['LITHOLOGY_GEOLINK'].astype('category').values
plot_well(well_name,
          logs,
          facies)
logs

def dset_to_nc(dset, f, engine="netcdf4", compression={"zlib": True}):
    if isinstance(dset, xr.DataArray):
        dset = dset.to_dataset(name="data")
    encoding = {k: {"zlib": True} for k in dset.data_vars}
    print('saving to {}'.format(f))
    dset.to_netcdf(f, engine=engine, encoding=encoding)
    print('Wrote {}.nc size={} M'.format(f.stem, f.stat().st_size / 1000000.0))
dset_to_nc(dset=xr_all.drop(['geometry']),
          f=interim_locations/'geolink_norge_well_logs.h5')
# +
import os, shutil

def get_dir_size(start_path="."):
    total_size = 0
    for dirpath, dirnames, filenames in os.walk(start_path):
        for f in filenames:
            fp = os.path.join(dirpath, f)
            total_size += os.path.getsize(fp)
    return total_size

def dset_to_zarr(dset, f):
    if isinstance(dset, xr.DataArray):
        dset = dset.to_dataset(name="data")
    encoding = {k: {"zlib": True} for k in dset.data_vars}
    print('saving to {}'.format(f))
    if f.exists():
        try:
            return xr.open_zarr(f)
        except:
            shutil.rmtree(f)
    dset.to_zarr(str(f))
    print('{}.zarr size={} M'.format(f.stem, get_dir_size(str(f)) / 1000000.0))

dset_to_zarr(dset=xr_all.drop(['geometry']),
          f=interim_locations/'geolink_norge_well_logs.zarr')
# -
# # Plot map

import matplotlib.pyplot as plt
# %matplotlib inline
import os
from tqdm.auto import tqdm
import pandas as pd
import geopandas as gpd
import numpy as np

# +
# import pandas as pd
# import xarray as xr
# xf = xr.open_zarr("../../data/processed/geolink_norge_dataset/geolink_norge_well_logs.zarr")
# df = xf.to_dataframe().swaplevel().sample(1000)
# df['LITHOLOGY_GEOLINK'] = df['LITHOLOGY_GEOLINK'].astype('category')
# df['Well'] = df.index.get_level_values(0).astype('category')
# df['DEPT'] = df.index.get_level_values(1)
# feature_cols = ['CALI', 'DTC', 'GR', 'RDEP', 'RHOB',
#        'RMED', 'xc', 'yc', 'DEPT']
# df = df.dropna(how='any', subset=feature_cols+['LITHOLOGY_GEOLINK'])
# df = df.sort_index()

# import geopandas as gpd
# gdf = gpd.GeoDataFrame(df, geometry=gpd.points_from_xy(df.xc, df.yc))
# gdf = gdf.set_crs(epsg=4326).to_crs(epsg=3857)
# gdf.plot()
# -

# # Plot contextily

from pathlib import Path
interim_locations = Path("../../data/processed/geolink_norge_dataset/")
df_well_tops = gpd.read_file(interim_locations / "norge_well_tops.gpkg").set_crs(epsg=4326).to_crs(epsg=3857)#.head(40)
# df_well_tops.plot()

# +

import contextily as ctx
ax = df_well_tops.plot(figsize=(18, 18), edgecolor='k')
ctx.add_basemap(ax, url=ctx.providers.Esri.OceanBasemap, zoom=8)

# Plot every 5th
df_well_tops[::5].apply(lambda x:
                   ax.annotate(
                       s=x.wlbWellboreName,
                       xy=x.geometry.centroid.coords[0],
                       ha='left',
                       c='white',

                   ), axis=1);
# +
ax = df_well_tops.plot(figsize=(18, 18), edgecolor='k')
# ctx.add_basemap(ax, url=ctx.providers.Esri.OceanBasemap)
ctx.add_basemap(ax,
                crs=df_well_tops.crs.to_string(),
                source=ctx.providers.Stamen.Watercolor
               )

# Plot every 5th
df_well_tops[::5].apply(lambda x:
                   ax.annotate(
                       s=x.wlbWellboreName,
                       xy=x.geometry.centroid.coords[0],
                       ha='left',
                       c='white'
                   ), axis=1);

# -

west, south, east, north = bbox = df_well_tops.total_bounds
img, ext = ctx.bounds2raster(west,
                             south,
                             east,
                             north,
                             "world_watercolor.tif",
                             source=ctx.providers.Stamen.Watercolor,
                             ll=True,
                             zoom=8
                            )

west, south, east, north = bbox = df_well_tops.total_bounds
img, ext = ctx.bounds2raster(west,
                             south,
                             east,
                             north,
                             "oceanesri.tif",
                             source=ctx.providers.Esri.OceanBasemap,
                             ll=True,
                             zoom=8
                            )

# +
# ctx.bounds2raster?
# -


# -*- encoding: utf-8 -*-
"""
@Author     :   zYx.Tom
@Contact    :   526614962@qq.com
@site       :   https://zhuyuanxiang.github.io
---------------------------
@Software   :   PyCharm
@Project    :   deep-learning-with-python-notebooks
@File       :   ch0702_keras_callbacks.py
@Version    :   v0.1
@Time       :   2019-11-27 11:54
@License    :   (C)Copyright 2018-2019, zYx.Tom
@Reference  :   《Python 深度学习，Francois Chollet》, Sec0702，P210
@Desc       :   高级的深度学习最佳实践，使用Keras回调函数
"""
import os
import sys

import keras
import matplotlib.pyplot as plt
import numpy as np  # pip install numpy<1.17，小于1.17就不会报错
import winsound
from keras.callbacks import EarlyStopping
from keras.callbacks import ModelCheckpoint
from keras.callbacks import ReduceLROnPlateau
from keras.losses import binary_crossentropy
from keras.models import Sequential
from keras.optimizers import rmsprop

# 屏蔽警告：Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
# 设置数据显示的精确度为小数点后3位
np.set_printoptions(precision = 3, suppress = True, threshold = np.inf, linewidth = 200)
# to make this notebook's output stable across runs
seed = 42
np.random.seed(seed)
# Python ≥3.5 is required
assert sys.version_info >= (3, 5)
# numpy 1.16.4 is required
assert np.__version__ in ["1.16.5", "1.16.4"]

# ----------------------------------------------------------------------
# 设置模型检查点，监控模型达到要求时就停止训练
def model_check_and_stop():
    callbacks_list = [
            # 监控模型的验证精度，如果精度在多于一轮时间（即两轮）内不再改善，就中断训练
            EarlyStopping(monitor = 'acc', patience = 1),
            # 每轮过后都保存当前权重于目标文件中，如果‘val_loss’没有改善，就不需要覆盖模型文件，目标是始终保存在训练过程中见到的最佳模型
            ModelCheckpoint(filepath = 'my_model.h5', monitor = 'val_loss', save_best_only = True)
    ]
    model = Sequential()
    model.compile(optimizer = rmsprop(), loss = binary_crossentropy, metrics = ['acc'])
    # 由于回调函数要监控“验证集损失”和“验证集精度”，所以需要传入validation_data（验证数据）
    model.fit(x, y, epochs = 10, batch_size = 32, callbacks = callbacks_list, validation_data = (x_val, y_val))

# ----------------------------------------------------------------------
# ReduceLROnPlateau（Reduce learning rate when a metric has stopped improving.）
# 当监控值到达一个稳定水平（Plateau），那么减少学习率
def reduce_lr_on_plateau():
    # 监控模型的“验证集损失”，如果验证损失在10轮内都没有改善，那么就触发这个回调函数，将学习率乘以0.1
    callbacks_list = [ReduceLROnPlateau(monitor = 'val_loss', factor = 0.1, patience = 10)]
    # 由于回调函数要监控“验证集损失”，所以需要传入validation_data（验证数据）
    model.fit(x, y, epochs = 10, batch_size = 32, callbacks = callbacks_list, validation_data = (x_val, y_val))

# ----------------------------------------------------------------------
# 编写自己的回调函数（创建keras.callbacks.Callback子类）
# 类的时间点：
# on_epoch_begin：在每轮开始时被调用
# on_epoch_end：在每轮结束时被调用
# on_batch_begin：在处理每个批量之前被调用
# on_batch_end：在处理每个批量之后被调用
# on_train_begin：在训练开始时被调用
# on_train_end：在训练结束时被调用
# 回调函数可以访问的属性：
# self.model：调用回调函数的模型实例
# self.validation_data：传入fit()作为验证数据的值
class MyCallBack(keras.callbacks.Callback):
    # 在每轮结束后将模型每层的激活保存到硬盘中（格式为 Numpy 数组）
    # 这个激活是对验证集的第一个样本计算得到的
    def set_model(self, model):
        self.model = model
        layer_outputs = [layer.output for layer in model.layers]
        self.activations_model = keras.models.Model(model.input, layer_outputs)
        pass

    def on_epoch_end(self, epoch, logs = None):
        if self.validation_data is None:
            raise RuntimeError("Requires validation_data.")
        # 获取验证数据的第一个输入样本
        validation_sample = self.validation_data[0][0:1]
        activations = self.activations_model.predict(validation_sample)
        f = open('activations_at_epoch_' + str(epoch) + '.npz', 'w')
        np.savez(f, activations)
        f.close()

# ----------------------------------------------------------------------
# 运行结束的提醒
winsound.Beep(600, 500)
if len(plt.get_fignums()) != 0:
    plt.show()
pass