# google driveをマウント
from google.colab import drive
drive.mount('/content/drive')

# torch_xla をインストール
DIST_BUCKET="gs://tpu-pytorch/wheels"
TORCH_WHEEL="torch-1.15-cp36-cp36m-linux_x86_64.whl"
TORCH_XLA_WHEEL="torch_xla-1.15-cp36-cp36m-linux_x86_64.whl"
TORCHVISION_WHEEL="torchvision-0.3.0-cp36-cp36m-linux_x86_64.whl"

!pip uninstall -y torch torchvision
!gsutil cp "$DIST_BUCKET/$TORCH_WHEEL" .
!gsutil cp "$DIST_BUCKET/$TORCH_XLA_WHEEL" .
!gsutil cp "$DIST_BUCKET/$TORCHVISION_WHEEL" .
!pip install "$TORCH_WHEEL" "$TORCH_XLA_WHEEL" "$TORCHVISION_WHEEL"
!sudo apt-get install libomp5

import torch
import torch.nn as nn
import torch.optim as optim
import torch.autograd as autograd
import torchvision.models as models
from torch.utils.data import Dataset, DataLoader
import torch_xla.distributed.data_parallel as dp
import torch_xla.core.xla_model as xm
from sklearn.model_selection import train_test_split
import numpy as np
import matplotlib.pyplot as plt
import albumentations
import random
import copy
import os
import csv

#  乱数設定
RANDOM_SEED = 2020
random.seed(RANDOM_SEED)
os.environ['PYTHONHASHSEED'] = str(RANDOM_SEED)
np.random.seed(RANDOM_SEED)
torch.manual_seed(RANDOM_SEED)

print(f'random seed={RANDOM_SEED}')

class UkiyoeDataset(Dataset):
  def __init__(self, images, labels, train=False):
    super().__init__()
    self.images = images
    self.labels = labels
    
    transforms = [albumentations.Normalize(
      mean=[0.721, 0.659, 0.560], std=[0.224, 0.207, 0.198], p=1)]
    
    if train:
      transforms.append(albumentations.HorizontalFlip(p=0.5))

    self.transform = albumentations.Compose(transforms)

  def __len__(self):
    return len(self.images)

  def __getitem__(self, idx):
    image = np.rollaxis(self.transform(image=self.images[idx])['image'], 2, 0)
    label = self.labels[idx]
    
    return image, label


# データローダの作成
train_labels = np.load(
  'drive/My Drive/ProbSpace/ukiyoe/ukiyoe-train-labels.npz')['arr_0']
train_images = np.load(
  'drive/My Drive/ProbSpace/ukiyoe/ukiyoe-train-imgs.npz')['arr_0']

train_x, valid_x, train_y, valid_y = train_test_split(
  train_images, train_labels, test_size=0.2, random_state=RANDOM_SEED)

train_loader = DataLoader(
  UkiyoeDataset(train_x, train_y, train=True),
  batch_size=64, shuffle=True)
valid_loader = DataLoader(
  UkiyoeDataset(valid_x, valid_y, train=False),
  batch_size=64, shuffle=False)

print(f'train data size: {len(train_x)}')
print(f'valid data size: {len(valid_x)}')

class Classifier(nn.Module):
  def __init__(self, model):
    super().__init__()
    self.model = model
    self.criterion = nn.CrossEntropyLoss()
  
  def forward(self, images, labels):
    preds = self.model(images)
    loss = self.criterion(preds, labels)
    accuracy = preds.max(dim=1)[1].eq(labels).float().mean()

    return loss, accuracy


def perform(model, loader, optimizer):
  loss_total = 0
  accuracy_total = 0
  count = 0
  
  for images, labels in loader:
    loss, accuracy = model(images, labels)
   
    if optimizer is not None:
      optimizer.zero_grad()
      loss.backward()
      xm.optimizer_step(optimizer)
    
    loss_total += float(loss) * len(images)
    accuracy_total += float(accuracy) * len(images)
    count += len(images)

  return np.array(
    (loss_total / count, accuracy_total / count), dtype=np.float64)


def train_func(model, loader, device, context):
  optimizer = context.getattr_or('optimizer',
    lambda: optim.Adam(model.parameters(), lr=0.001))
  
  model.train()
  
  return perform(model, loader, optimizer)


def valid_func(model, loader, device, context):
  model.eval()
  
  return perform(model, loader, None)


# モデル作成
model_parallel = dp.DataParallel(Classifier(
    models.resnet18(pretrained=False, num_classes=10)))
print(f'tpu devices: {model_parallel.devices}')

# 学習
log = []

for epoch in range(10):
  results = model_parallel(train_func, train_loader)
  train_loss, train_accuracy = sum(results) / len(results)
  
  results = model_parallel(valid_func, valid_loader)
  valid_loss, valid_accuracy = sum(results) / len(results)
  
  print('[{}] train(loss/accuracy)={:.2f}/{:.2f}, valid(loss/accuracy)={:.2f}/{:.2f}'.format(
    epoch + 1, train_loss, train_accuracy, valid_loss, valid_accuracy))
  
  log.append((epoch + 1, train_loss, train_accuracy, valid_loss, valid_accuracy))

# モデルを保存
model_copy = copy.deepcopy(model_parallel.models[0].model).to('cpu')
torch.save(model_copy.state_dict(), 'resnet18.pth')
  
# 結果表示
figure = plt.figure(figsize=(8, 3))

axis = figure.add_subplot(1, 2, 1)
axis.plot([x[0] for x in log], [x[1] for x in log], label='train')
axis.plot([x[0] for x in log], [x[3] for x in log], label='valid')
axis.set_xlabel('epoch')
axis.set_ylabel('loss')
axis.legend()

axis = figure.add_subplot(1, 2, 2)
axis.plot([x[0] for x in log], [x[2] for x in log], label='train')
axis.plot([x[0] for x in log], [x[4] for x in log], label='valid')
axis.set_xlabel('epoch')
axis.set_ylabel('accuracy')
axis.legend()

plt.show()

# モデル読み込み
model = models.resnet18(pretrained=False, num_classes=10)
model.load_state_dict(torch.load('resnet18.pth', map_location=lambda s, _: s))

# 推論
results = [[0] * 10 for _ in range(10)]

device = xm.xla_device()
model = model.to(device)
model.eval()

for images, labels in valid_loader:
  with torch.no_grad():
    preds = model(images.to(device)).max(dim=1)[1].cpu()
 
  for label, pred in zip(labels, preds):
    results[label][pred] += 1

true_pos_total = 0
false_pos_total = 0
false_neg_total = 0
fscore_total = 0
lbscore_total = 0

# 表示
test_dists = [ # テストデータに含まれる各クラスの割合
    0.181, 0.209, 0.093, 0.068, 0.113, 0.063, 0.116, 0.048, 0.058, 0.050]

print('-- predictions (L:Labels, P:Predictions)')
print('   |{}'.format('|'.join(f'P={i}' for i in range(10))))
for i in range(10):
  print('L={}|{}'.format(i, '|'.join(f'{r:3d}' for r in results[i])))

print()
print('-- scores')
print('   |count|precision|recall|f-score')
for i in range(10):
  true_pos = results[i][i]
  false_neg = sum(results[i]) - true_pos
  false_pos = sum(results[j][i] for j in range(10)) - true_pos
  precision = true_pos / (true_pos + false_pos)
  recall = true_pos / (true_pos + false_neg)
  fscore = (2 * precision * recall) / (precision + recall)

  true_pos_total += true_pos
  false_pos_total += false_pos
  false_neg_total += false_neg
  fscore_total += fscore
  lbscore_total += test_dists[i] * recall

  print('L={}|  {:3d}|   {:.4f}|{:.4f}| {:.4f}'.format(
    i, sum(results[i]), precision, recall, fscore))

print()
print('accuracy       = {:.4f}'.format(true_pos_total / len(valid_x)))
print('f-score(micro) = {:.4f}'.format(
  (2 * true_pos_total) / (2 * true_pos_total + false_pos_total + false_neg_total)))
print('f-score(macro) = {:.4f}'.format(fscore_total / 10))
print('estimated LB   = {:.4f}'.format(lbscore_total))

# データ読み込み
test_images = np.load(
    'drive/My Drive/ProbSpace/ukiyoe/ukiyoe-test-imgs.npz')['arr_0']
test_labels = np.zeros((test_images.shape[0], 1), dtype=np.int64)

test_loader = DataLoader(
  UkiyoeDataset(test_images, test_labels, train=False),
  batch_size=64, shuffle=False)

# 推論
predictions = []
model.eval()

for images, _ in test_loader:
  with torch.no_grad():
    preds = model(images.to(device)).max(dim=1)[1].cpu()
  
  predictions.extend(int(x) for x in preds)

# results
classes = [0] * 10
dists = [int(round(d * len(predictions))) for d in test_dists]
error = 0

for p in predictions:
  classes[p] += 1

print("-- number of predictions")
for i, (c, d) in enumerate(zip(classes, dists)):
  error += abs(c - d)
  print(f'class-{i}: {c:3d} (actual={d})')
  
# 提出ファイル作成
with open('submission.csv', 'w') as writer:
  csv_writer = csv.writer(writer)
  csv_writer.writerow(('id', 'y'))
  csv_writer.writerows((i + 1, x) for i, x in enumerate(predictions))