import matplotlib.pyplot as plt
import numpy as np

from . import utils

class Resonance():
	"""Stores chemical shifts for an atom and analytical functions"""

	def __init__(self, key, shifts):
		"""initialise resonance

		:param key: Descriptive key for atom 'chainID.resName.resID.name'
		:type key: str
		:param shifts: Dictionary of chemical shifts (float) indexed by 
			frame (int)
		:type shifts: dict
		"""
		info = key.split('.')
		self.chainID = int(info[0])
		self.resName = info[1]
		self.resSeq = int(info[2])
		self.name = info[3]
		self.shifts = shifts
		

	def average(self):
		"""Get to average chemical shift of resonance

		:return: Average chemical shift
		:rtype: float
		"""		
		# Order frames and get shifts
		sorted_frames = sorted(self.shifts.keys())
		shifts = [self.shifts[frame] for frame in sorted_frames]
		return np.mean(shifts)

	# Todo list
	def converge():
		pass


	def get_shifts(self):
		"""Return list of chemical shifts for resonance

		:return: List of tuples (frame, shift)
		:rtype: list of tuples (int, float)
		"""		
		sorted_frames = sorted(self.shifts.keys())
		return [(frame, self.shifts[frame]) for frame in sorted_frames]


	def plot(self, xlim=[], ylim=[], figsize=(8.0,4.0), textsize=16):
		"""Plot chemical shifts and running average for resonance

		:param xlim: X minimum and maximum [xmin,xmax], 
			defaults to [] (autogenerated)
		:type xlim: list of two floats, optional
		:param ylim: Y minimum and maximum [ymin,ymax]
		:type ylim: list of two floats, optional
		:param figsize: X and Y figure size in inches, defaults to (8.0,4.0)
		:type figsize: tuple of two floats, optional
		:param textsize: Text size of labels, defaults to 16
		:type textsize: int, optional
		"""
		# Order frames and get shifts
		sorted_frames = sorted(self.shifts.keys())
		shifts = [self.shifts[frame] for frame in sorted_frames]
		stats = utils.evolving_stats(shifts)
		means = [mean for mean, stdev in stats]

		# Get slope of final 20% for means
		fp = int(len(sorted_frames)*0.2) * -1
		# Only do if more than 3 points
		#print(fp)
		m, c = utils.fit_linear(sorted_frames[fp:], means[fp:])
		fit = [utils.linear(x, m, c) for x in sorted_frames]
		print(fp*m*-1)

		#means_plus_stdev = [mean + stdev for mean, stdev in stats]
		#means_minus_stdev = [mean - stdev for mean, stdev in stats]

		# Plot line plot
		fig, ax = plt.subplots(figsize=figsize)
		plt.xticks(fontsize=textsize)
		plt.yticks(fontsize=textsize)
		ax.set_xlabel('Frame', size=textsize)
		ax.set_ylabel('Chemical shift (ppm)', size=textsize)
		if xlim:
			ax.set_xlim(*xlim)
		if ylim:
			ax.set_ylim(*ylim)
		ax.scatter(sorted_frames, shifts, c='black')
		ax.plot(sorted_frames,means, c='black')
		ax.plot(sorted_frames,fit, c='black')
		#ax.plot(sorted_frames,means_plus_stdev, c='red')
		#ax.plot(sorted_frames,means_minus_stdev, c='red')
		plt.show()
