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

"""
Explanation: Copyright 2020 Google LLC.
Licensed under the Apache License, Version 2.0 (the "License")
End of explanation
"""


import pandas as pd
import functools
import bokeh.io
!pip install pandas_bokeh
import pandas_bokeh
import requests
import json
import numpy as np
bokeh.io.output_notebook()

"""
Explanation: Human evaluation of visual metrics
This colab explores correlations between the mucped22 dataset and various objective visual metrics.
Unlike many other datasets, these evaluations:
 *   Are made using only compression distortions (since the distortions were created using AVIF, JPEG, and JXL encoders), which will benefit metrics designed for compression artifacts.
 *   Are made by human evaluators experienced with image quality comparisons, which will benefit smaller distortions, relatively speaking, since unexperienced evaluators often don't notice them.
 *   Are made using two-alternative-forced-choice with a reference image instead of MOS, which will reduce the noise, since evaluators don't need to calibrate their MOS scores.
For each image, each distortion (method and quality setting) has then been ranked using ELO, to provide an expected human-rated ranking for each distortion.
This ranking will allow a comparison of the various metrics across different levels of distortion, e.g. near just-noticeable-differences vs far from just-noticeable-differences.
End of explanation
"""


!wget --quiet --no-check-certificate https://storage.googleapis.com/gresearch/mucped22/evaluations.json

with open('evaluations.json') as f:
  data = pd.DataFrame(json.load(f))

data

"""
Explanation: First download the dataset containing all evaluations.
End of explanation
"""


data['complete_crop'] = data.apply(lambda row: row.crop[0] + row.crop[2] <= row.image_dims[0] and row.crop[1] + row.crop[3] <= row.image_dims[1], axis=1)
data['worst_elo'] = data.apply(lambda row: row.greater_elo if row.greater_elo > row.lesser_elo else row.lesser_elo, axis=1)
data = data[(data.rater_flips > 2) & (data.rater_time_ms > 3000) & (data.complete_crop == True)]
data

def strip(ary, n):
  def stripfun(sum, el):
    sum[el] = el[n:]
    return sum
  return functools.reduce(stripfun, ary, {})

greater_metric_cols = list(filter(lambda el: el.startswith('greater_') and not el.endswith('_file'), list(data.columns)))

lesser_metric_cols = list(filter(lambda el: el.startswith('lesser_') and not el.endswith('_file'), list(data.columns)))

greater_metrics = data[greater_metric_cols]
greater_metrics = greater_metrics.rename(columns=strip(greater_metric_cols, 8))

lesser_metrics = data[lesser_metric_cols]
lesser_metrics = lesser_metrics.rename(columns=strip(lesser_metric_cols, 7))

"""
Explanation: Then decorate it with whether the crop settings were actually compatible with the image size (a few, ~15, evaluations have this bug), and the worst ELO of both distortions.
Finally filter out all evaluations where the evaluator didn't seem to do a good job (didn't flip between distortions more than 2 times, didn't spend more than 3 seconds on the evaluation).
End of explanation
"""


metrics = pd.concat([greater_metrics, lesser_metrics])
metrics = metrics.sort_values('elo').reset_index(drop=True)
metrics

"""
Explanation: To allow a rank correlation, like Spearman, combine the metrics of the worse distortion (lesser), and the better distortion (greater), into one dataframe. To also allow comparing correlation in different regions of quality, sort by ELO score.
End of explanation
"""


corrs = metrics.corr(method='spearman')
corrs

metric_cols = list(map(lambda name: name[7:], lesser_metric_cols))
metric_cols.remove('elo')

def rollingcorr(df, method, window_size, step_size):
  res = []
  for start in range(0, df.shape[0] - window_size, step_size):
    window = df[start:start+window_size]
    row = [window.iloc[-1]['elo']]
    for metric_name in metric_cols:
      row.append(np.abs(window[metric_name].corr(window['elo'], method=method)))
    res.append(row)
  return pd.DataFrame(res, dtype=np.float, columns=['elo'] + list(map(lambda name: f"{name}", metric_cols)))

"""
Explanation: Then compute the correlation matrix for these, using Spearman's rank correlation coeffient.
End of explanation
"""


rollingcorr(metrics, 'spearman', 5000, 1000).plot_bokeh(x='elo', figsize=(1400, 400))

"""
Explanation: Plot the correlation in a rolling window of 5000 evaluations with a step of 1000 evaluations for each metric, to see how they behave across a range of ELO scores.
End of explanation
"""


