%pylab inline

def meta_collect(ann_path: str, result_file: str, sep: str='\t') -> None:
    
    '''collect metadata for all images to "result_file"
    from json files in "ann_path" (execution time: about 5 mins)'''

    start = time()
    with open(result_file, 'w',  encoding='utf-8') as f:
        f.write(sep.join(['width', 'height', 'description',
                           'isModerated', 'moderatedBy', 'predicted']) + '\n')

    
        for file in tqdm(glob.glob(os.path.join(ann_path, '*.json'))):

            with open(file, encoding='utf-8') as js:
                tmp = json.load(js)

            try:
                f.write(sep.join([tmp['name'], str(tmp['size']['width']), str(tmp['size']['height']),
                               tmp['description'], str(tmp['moderation']['isModerated']),
                               tmp['moderation']['moderatedBy'], str(tmp['moderation']['predicted'])]) + '\n')
            except Exception:
                print(tmp['description'])
    print('execution time:', (time() - start), 'secs')

# meta_collect(ann_path, os.path.join(meta_path, 'metadata.tsv'))
df = pd.read_csv(os.path.join(meta_path, 'metadata.tsv'), sep='\t', index_col=0)

print(df.shape)
df.head()

df['description_length'] = df.description.apply(len)
df

df.describe().round(2)

df.describe(include=object)

df[~df.predicted.isna()]

df.isModerated.value_counts()

cou = df.moderatedBy.value_counts()
cou

n = 10
img_names = random.choice(df.index, n)
fig, axes = subplots(n, 1)

for img_name, ax in zip(img_names, axes):
    read_image(img_name + '.jpg', ax=ax, title=img_name + f'  ({df.loc[img_name].description})')
tight_layout()

# Creating dataframe with symbol counts with indexes from original df
counts = counts_to_df(df, 'description') # Do not dropping duplicates for saving original df indexes

# Barplot with symbol counts in dataset
fig, ax = subplots(figsize=(20, 10))
ax.set_title('Frequency of symbols')
color = '#2ca3db'
ax = sns.barplot(data=counts.sort_values('counts', ascending=False),
                 x='symbols', y='counts', ax=ax, color=color)
ax.set(xlabel='letter', ylabel='frequency')

tight_layout()

print(f'All characters:\n {counts.symbols.unique()}')

# Creating reference alphabet with Russian (lower- and uppercase) and punctuation symbols
alphabet_lower = [chr(ord("а") + i) for i in range(32)] + [chr(ord("а") + 33)] # Last is "ё"
alphabet_upper = [chr(ord("А") + i) for i in range(32)]
punctuation = list(string.punctuation)
ref_alphabet = set(alphabet_lower + alphabet_upper + punctuation)

# Creating alphabet from dataset
alphabet = set(counts.symbols)

print(f'Reference alphabet length: {len(ref_alphabet)}')
print(f'Actual alphabet length: {len(alphabet)}\n')

# difference between dataset and reference alphabet
diff_symbols = alphabet - ref_alphabet 
diff_counts = counts[counts.symbols.isin(diff_symbols)]
print(f'Non-ordinary symbols: {diff_symbols}')
print(f'Number of labels with non-ordinary symbols: {diff_counts.index.unique().shape[0]}')
print(f'Percentage of labels with non-ordinary symbols: {round(diff_counts.index.unique().shape[0]/df.shape[0]*100, 2)}%')

fig, axes = subplots(len(diff_symbols), 1)

for sym, ax in zip(diff_symbols, axes):
    ind = counts[counts.symbols == sym].index[0]
    read_image(ind + '.jpg', ax, df.loc[ind].description + f'  ({sym})')
tight_layout()

pd.options.display.max_rows = 100
df.loc[counts[counts.symbols.isin(diff_symbols)].index.drop_duplicates()].drop_duplicates('description')

counts = counts_to_df(df, 'description')

counts_dict = counts.set_index('symbols')['counts'].to_dict()
kazakh_symbols = set(counts_dict) - ref_alphabet
kazakh_counts = counts[counts.symbols.isin(kazakh_symbols)]
print(f'Kazakh symbols: {kazakh_symbols}')
print(f'Number of Kazakh labels: {kazakh_counts.index.unique().shape[0]}')
print(f'Percentage of Kazakh labels: {round(kazakh_counts.index.unique().shape[0]/df.shape[0]*100, 2)}%')

df = df.drop(counts[counts.symbols.isin(kazakh_symbols)].index.drop_duplicates(), axis=0)
df

fig, axes = subplots(2, 1, figsize=(15, 15))

tmp = counts_to_df(df)
tmp = tmp[~tmp.symbols.isin(punctuation)].drop_duplicates('symbols')

dictir = {}
for i, j in tmp.iterrows():
    sym, cou = j
    dictir[sym.lower()] = dictir.get(sym.lower(), 0) + cou
tmp = pd.DataFrame(dictir, index=['counts']
                  ).T.reset_index().sort_values('counts', ascending=False)

sns.barplot(data=tmp, x='index', y='counts', ax=axes[0], color=color)
axes[0].set(xlabel='letter', ylabel='frequency')
axes[0].set_title('Dataset alphabet')

ls = pd.read_csv(os.path.join('..', 'metadata', 'alphabet.tsv'), sep='\t', 
                 index_col=1).sort_values('Частотность', ascending=False)
sns.barplot(data=ls, x='Буква', y='Частотность', ax=axes[1], color=color)
axes[1].set(xlabel='letter', ylabel='frequency')
axes[1].set_title('Russian alphabet')
tight_layout()

plt.figure(figsize=(15, 8))
ax = sns.kdeplot(df.description_length, bw=0.3)
ax.set(xlabel='description length', ylabel='density', title='Length distribution of text');

print(f'Color format: RGB')
print_imgs_with_hist(df[['description']])

plt.figure(figsize=(15, 8))
ax = sns.scatterplot(x=df.width, y=df.height, alpha=0.4);
ax.set(xlabel='description length', ylabel='density', title='Size distribution of pictures');

print(f"width: max = {df.width.max()}, min = {df.width.min()}, mean = {df.width.mean()}")
print(f"height: max = {df.height.max()}, min = {df.height.min()}, mean = {df.height.mean()}\n")

mheight = df.height.value_counts().sort_values(ascending=False)
mwidth = df.width.value_counts().sort_values(ascending=False)

print(f'most common widths (of {mwidth.shape[0]} size):\nwidth\tcount\n{mwidth.head(10)}\nand their mean = {mwidth.head(10).index.to_series().mean()}\n')
print(f'most common heights (of {mheight.shape[0]} size):\nheight\tcount\n{mheight.head(10)}\nand their mean = {mheight.head(10).index.to_series().mean()}')

fig, axes = subplots(4, 1, figsize=(10, 15))

df_list = [df[df.height == mheight.index[0]],
           df[df.height == mheight.index[-1]],
           df[df.width == mwidth.index[0]],
           df[df.width == mwidth.index[-1]]]

titles = ['most common height example',
          'least common height example',
          'most common width example',
          'least common width example']

for tmp, title, ax in zip(df_list, titles, axes):
    read_image(tmp.index[0] + '.jpg', ax, title=title + f' ({tmp.description[0]})')
    ax.axis('on')

tight_layout()

fig, axes = subplots(6, 1, figsize=(10, 15))

df_list = [df[df.height == df.height.max()],
           df[df.height == df.height.min()],
           df[df.width == df.width.max()],
           df[df.width == df.width.min()],
           df[df.description.apply(len) == df.description.apply(len).max()],
           df[df.description.apply(len) == df.description.apply(len).min()]]

titles = ['max hight example',
          'min hight example',
          'max width example',
          'min width example',
          'max len description',
          'min len description']

for tmp, title, ax in zip(df_list, titles, axes):
    read_image(tmp.index[0] + '.jpg', ax, title=title + f' ({tmp.description[0]})')
    ax.axis('on')

tight_layout()

img_height, img_width = 100, 600
df = df[(df.width <= img_width) & (df.height <= img_height)]
max_length = df.description.str.len().max()
print(max_length)
df

# There is 1.0 in cell if there is this symbol in this picture otherwise 0 
counts = counts_to_df(df, 'description')
counts.counts = 1
splitter = counts.reset_index().drop_duplicates().pivot(index='index', columns='symbols').fillna(0)
splitter

# treat splitter df as multilabel class signature, so we can easy split original df to train and test
train, test, _, ls = train_test_split(df, splitter, shuffle=True,
                            test_size=0.15, random_state=12)

# And then split test to final test and val dfs
test, val, _, _ = train_test_split(test, ls, shuffle=True,
                            test_size=0.33, random_state=17)

train_counts = counts_to_df(train, 'description')
test_counts = counts_to_df(test, 'description')
val_counts = counts_to_df(val, 'description')

print('Sets differences between presented symbols in train, test and val data\n')

print('train_counts - test_counts:   ', set(train_counts.symbols) - set(test_counts.symbols))
print('train_counts - val_counts:   ', set(train_counts.symbols) - set(val_counts.symbols))

print('test_counts - train_counts:   ', set(test_counts.symbols) - set(train_counts.symbols))
print('test_counts - val_counts:   ', set(test_counts.symbols) - set(val_counts.symbols))

print('val_counts - train_counts:   ', set(val_counts.symbols) - set(train_counts.symbols))
print('val_counts - test_counts:   ', set(val_counts.symbols) - set(test_counts.symbols))

# Plot frequencies of symbols in three new dataframes
fig, axes = subplots(3, 1, figsize=(15, 15))

for tmp, ax, name in zip((train_counts, test_counts, val_counts), axes, ['train', 'test', 'val']):
    sns.barplot(data=tmp.sort_values('counts', ascending=False),
                x='symbols', y='counts', ax=ax)
    ax.set_title(name)
tight_layout()

# Mapping characters to integers
counts = counts_to_df(df)
counts = counts[~counts.isin(['', ' '])].symbols.unique().tolist() + [' ', '#']
vocab = pd.Series(counts).str.encode('utf8')

char_to_num = layers.experimental.preprocessing.StringLookup(
    vocabulary=vocab,
    mask_token=None,
)

# Mapping integers back to original characters
num_to_char = layers.experimental.preprocessing.StringLookup(
    vocabulary=char_to_num.get_vocabulary(), mask_token=None, invert=True,
)

blank_index = char_to_num(tf.strings.unicode_split('#', input_encoding="UTF-8")).numpy()[0]
blank_index  # For blank symbols

ind = df.index.tolist()
# random.shuffle(ind)
ls = get_dataset(df.loc[ind].iloc[:16])
for batch in ls.take(1):
    show_batch(batch, batch_size=batch_size)

paths = df.index.to_series().apply(lambda x: os.path.join(img_path, x) + '.jpg')

aug_1 = os.path.join(img_path, 'aug_1')

if not os.path.exists(aug_1):
    os.mkdir(aug_1)
sometimes = lambda aug: iaa.Sometimes(0.5, aug)
seq = iaa.Sequential(
    [

        iaa.Sometimes(0.1, iaa.GaussianBlur(3.0)),
    
        iaa.Sometimes(0.1, iaa.AveragePooling(2)),
        iaa.Sometimes(0.1, iaa.Emboss(alpha=(0.0, 1.0), strength=(0.75, 1.25))),
        iaa.Sometimes(0.1, iaa.GammaContrast((0.5, 1.0))),
        iaa.Invert(0.05, per_channel=True),
        iaa.Sometimes(0.1, iaa.CoarseDropout((0.0, 0.05), size_percent=(0.02, 0.25))),

        iaa.ElasticTransformation(alpha=(0.5, 3.5), sigma=0.25),

        iaa.PerspectiveTransform(scale=(0.02, 0.05)),

        iaa.Sometimes(0.1, iaa.SaltAndPepper(0.05)),
    ],
    random_order=True
)

for path in tqdm(paths):
    break
    print('start')
    img = imageio.imread(path)
    image = [copy(img) for _ in range(30)]

    ls = seq(images=image)
    
    print(len(ls))
    for i in range(30):
        _, name = os.path.split(path)
        name = os.path.join(aug_1, f'{i}_aug_' + name)
        cv2.imwrite(name, ls[i])

from preprocess import PreprocessFrame, Dataset, make_augments

img_width = 600
img_height = 100

# parameters of resized images
new_img_width = 300
new_img_height = 50

# default paths
WORKING_DIR = os.path.join('/home', 'htr')
ann_path = os.path.join(WORKING_DIR, 'HKR_Dataset_Words_Public', 'ann')
img_path = os.path.join(WORKING_DIR, 'HKR_Dataset_Words_Public', 'img')

# collect metadata
# meta_collect(ann_path, os.path.join(WORKING_DIR, 'metadata', 'metadata.tsv'))

# get preprocessed metadata dataframe
df = PreprocessFrame(metadata=os.path.join(WORKING_DIR, 'metadata', 'metadata.tsv'),
                     img_height=img_height, img_width=img_width)
print(df.shape)

# Make augments file (if they exists: comment or delete line)
aug_df = None
# aug_df = make_augments(df=df, img_path=img_path, WORKING_DIR=WORKING_DIR,
#                         img_height=img_height, img_width=img_width)

# get augments metadata dataframe from original dataframe if not starting make_augments
if not isinstance(aug_df, pd.DataFrame):
    aug_df = df.copy()
    aug_df.index = aug_df.index.to_series().apply(lambda x: os.path.join('aug_1', 'aug_' + x))

train, test, val = list(Dataset(df, aug_df=aug_df,
                                test_size=0.1,
                                val_size=0.05,
                                img_path=img_path,
                                img_height=img_height,
                                img_width=img_width,
                                new_img_height=new_img_height,
                                new_img_width=new_img_width,
                                WORKING_DIR=WORKING_DIR,
                                shuffle=True,
                                random_state=12))
train

for batch in train.take(1):
    show_batch(batch, batch_size=batch_size)

for batch in test.take(1):
    show_batch(batch, batch_size=batch_size)

for batch in val.take(1):
    show_batch(batch, batch_size=batch_size)