<commit_before># ex:ts=4:sw=4:sts=4:et
# -*- tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*-
import copy
import datetime
import hashlib
import json
import logging
import re
import time
from urllib.parse import parse_qs
from urllib.parse import urljoin
from urllib.parse import urlparse

from svtplay_dl.error import ServiceError
from svtplay_dl.fetcher.dash import dashparse
from svtplay_dl.fetcher.hls import hlsparse
from svtplay_dl.service import MetadataThumbMixin
from svtplay_dl.service import Service
from svtplay_dl.subtitle import subtitle
from svtplay_dl.utils.text import filenamify

URL_VIDEO_API = "https://api.svt.se/video/"
LIVE_CHANNELS = {
    "svtbarn": "ch-barnkanalen",
    "svt1": "ch-svt1",
    "svt2": "ch-svt2",
    "svt24": "ch-svt24",
    "kunskapskanalen": "ch-kunskapskanalen",
}

class Svtplay(Service, MetadataThumbMixin):
    supported_domains = ["svtplay.se", "svt.se", "beta.svtplay.se", "svtflow.se"]
    info_search_expr = r"<script id=\"__NEXT_DATA__\" type=\"application\/json\">({.+})<\/script>"
    access = None

    def get(self):
        parse = urlparse(self.url)
        if parse.netloc in ("www.svtplay.se", "svtplay.se"):
            if parse.path[:6] != "/video" and parse.path[:6] != "/klipp" and parse.path[:8] != "/kanaler":
                yield ServiceError("This mode is not supported anymore. Need the url with the video.")
                return

        query = parse_qs(parse.query)

        if "accessService" in query:
            self.access = query["accessService"]

        urldata = self.get_urldata()

        if parse.path[:8] == "/kanaler":
            ch = LIVE_CHANNELS[parse.path[parse.path.rfind("/") + 1 :]]
            _url = urljoin(URL_VIDEO_API, ch)
            res = self.http.get(_url)
            try:
                janson = res.json()
            except json.decoder.JSONDecodeError:
                yield ServiceError(f"Can't decode api request: {res.request.url}")
                return
            videos = self._get_video(janson)
            self.config.set("live", True)
            yield from videos
            return

        match = re.search(self.info_search_expr, urldata)
        if not match:
            yield ServiceError("Can't find video info.")
            return

        janson = json.loads(match.group(1))
        video_data = None
        vid = None
        for data_entry in janson["props"]["urqlState"].values():
            if "data" in data_entry:
                entry = json.loads(data_entry["data"])
                for key, data in entry.items():
                    if key == "detailsPageByPath" and data and "moreDetails" in data:
                        video_data = data
                        vid = data["video"]["svtId"]
                        break

        if not vid:
            yield ServiceError("Can't find video id")
            return

        self.outputfilename(video_data)
        self.extrametadata(video_data)

        res = self.http.get(URL_VIDEO_API + vid)
        try:
            janson = res.json()
        except json.decoder.JSONDecodeError:
            yield ServiceError(f"Can't decode api request: {res.request.url}")
            return
        if res.status_code >= 400:
            yield ServiceError("Can't find any videos. its removed?")
            return
        videos = self._get_video(janson)
        yield from videos

    def _get_video(self, janson):
        if "subtitleReferences" in janson:
            for i in janson["subtitleReferences"]:
                if i["format"] == "webvtt" and "url" in i:
                    yield subtitle(copy.copy(self.config), "wrst", i["url"], "sv", output=self.output)

        if "videoReferences" in janson:
            if len(janson["videoReferences"]) == 0:
                yield ServiceError("Media doesn't have any associated videos.")
                return

            for i in janson["videoReferences"]:
                if i["url"].find(".m3u8") > 0:
                    yield from hlsparse(self.config, self.http.request("get", i["url"]), i["url"], self.output)
                elif i["url"].find(".mpd") > 0:
                    yield from dashparse(self.config, self.http.request("get", i["url"]), i["url"], output=self.output)

    def _last_chance(self):
        videos = []
        match = re.search(self.info_search_expr, self.get_urldata())
        if not match:
            logging.error("Can't find video info.")
            return videos
        janson = json.loads(match.group(1))
        video_data = None
        for data_entry in janson["props"]["urqlState"].values():
            entry = json.loads(data_entry["data"])
            for key, data in entry.items():
                if key == "startForSvtPlay":
                    video_data = data
        if not video_data:
            return videos
        for section in video_data["selections"]:
            for i in section["items"]:
                videos.append(urljoin("http://www.svtplay.se", i["item"]["urls"]["svtplay"]))
        return videos

    def _genre(self):
        videos = []
        episodes = []
        singles = []
        parse = urlparse(self.url)
        url = f"https://www.svtplay.se{parse.path}?tab=all"
        data = self.http.get(url).text
        match = re.search(self.info_search_expr, data)
        if not match:
            logging.error("Can't find video info.")
            return episodes
        janson = json.loads(match.group(1))
        video_data = None
        for data_entry in janson["props"]["urqlState"].values():
            entry = json.loads(data_entry["data"])
            for key, data in entry.items():
                if key == "categoryPage":
                    if "lazyLoadedTabs" in data:
                        video_data = data["lazyLoadedTabs"]
        if not video_data:
            return episodes
        for lazytab in video_data:
            if "selections" in lazytab:
                for section in lazytab["selections"]:
                    for i in section["items"]:
                        if i["item"]["__typename"] == "Single":
                            singles.append(urljoin("http://www.svtplay.se", i["item"]["urls"]["svtplay"]))
                        else:
                            videos.append(urljoin("http://www.svtplay.se", i["item"]["urls"]["svtplay"]))
        for i in videos:
            episodes.extend(self._all_episodes(i))
        if singles:
            episodes.extend(singles)
        return episodes

    def _all_episodes(self, url):
        parse = urlparse(url)
        tab = None
        videos = []
        if parse.query:
            query = parse_qs(parse.query)
            if "tab" in query:
                tab = query["tab"][0]

        data = self.http.get(url).text
        match = re.search(self.info_search_expr, data)
        if not match:
            logging.error("Can't find video info.")
            return videos

        janson = json.loads(match.group(1))
        video_data = None
        for data_entry in janson["props"]["urqlState"].values():
            if "data" in data_entry:
                entry = json.loads(data_entry["data"])
                # logging.info(json.dumps(entry))
                for key, data in entry.items():
                    if key == "detailsPageByPath" and data and "heading" in data:
                        video_data = data
                        break

        collections = []
        if video_data is None:
            return videos
        if video_data["item"]["parent"]["__typename"] == "Single":
            videos.append(urljoin("http://www.svtplay.se", video_data["item"]["urls"]["svtplay"]))
        for i in video_data["associatedContent"]:
            if tab:
                if tab == i["id"]:
                    collections.append(i)
            else:
                if i["id"] == "upcoming" or i["id"] == "related":
                    continue
                elif self.config.get("include_clips") and "clips" in i["id"]:
                    collections.append(i)
                elif "clips" not in i["id"]:
                    collections.append(i)

        for i in collections:
            for epi in i["items"]:
                if epi["item"]["urls"]["svtplay"] not in videos:
                    videos.append(urljoin("http://www.svtplay.se", epi["item"]["urls"]["svtplay"]))
        return videos

    def find_all_episodes(self, config):
        parse = urlparse(self._url)

        episodes = []
        if re.search("^/sista-chansen", parse.path):
            episodes = self._last_chance()
        elif re.search("^/genre", parse.path):
            episodes = self._genre()
        else:
            episodes = self._all_episodes(self.url)

        if not episodes:
            logging.error("Can't find any videos.")
        else:
            if config.get("all_last") > 0:
                return episodes[-config.get("all_last") :]
        return episodes

    def outputfilename(self, data):
        name = None
        desc = None

        name = data["moreDetails"]["heading"]
        other = data["moreDetails"]["episodeHeading"]
        vid = hashlib.sha256(data["video"]["svtId"].encode("utf-8")).hexdigest()[:7]

        if name == other:
            other = None
        elif name is None:
            name = other
            other = None
        elif other is None:
            if name != data["moreDetails"]["titleHeading"]:
                other = data["moreDetails"]["titleHeading"]

        season, episode = self.seasoninfo(data)
        if "accessibility" in data:
            if data["accessibility"] == "AudioDescribed":
                desc = "syntolkat"
            if data["accessibility"] == "SignInterpreted":
                desc = "teckentolkat"

        if not other:
            other = desc
        elif desc:
            other += f"-{desc}"

        self.output["title"] = filenamify(name)
        self.output["id"] = vid
        self.output["season"] = season
        self.output["episode"] = episode
        self.output["episodename"] = other

    def seasoninfo(self, data):
        season, episode = None, None

        season_nr = self._find_season(data)
        episode_nr = self._find_episode(data)

        if season_nr is None or episode_nr is None:
            return season, episode

        season = f"{int(season_nr):02d}"
        episode = f"{int(episode_nr):02d}"

        return season, episode

    def _find_season(self, data):
        match = re.search(r"/säsong (\d+)/", data["analyticsIdentifiers"]["viewId"])
        if match:
            return match.group(1)

        match = re.search(r"-sasong-(\d+)-", data["item"]["urls"]["svtplay"])
        if match:
            return match.group(1)

        vid = data["video"]["svtId"]
        for seasons in data["associatedContent"]:
            for i in seasons["items"]:
                if i["item"]["videoSvtId"] == vid and "positionInSeason" in i["item"]:
                    match = re.search(r"S.song (\d+)", i["item"]["positionInSeason"])
                    if match:
                        return match.group(1)
        return None

    def _find_episode(self, data):
        match = re.search(r"/avsnitt (\d+)", data["analyticsIdentifiers"]["viewId"])
        if match:
            return match.group(1)

        match = re.search(r"Avsnitt (\d+)", data["item"]["name"])
        if match:
            return match.group(1)

        vid = data["video"]["svtId"]
        for seasons in data["associatedContent"]:
            for i in seasons["items"]:
                if i["item"]["videoSvtId"] == vid and "positionInSeason" in i["item"]:
                    match = re.search(r"Avsnitt (\d+)", i["item"]["positionInSeason"])
                    if match:
                        return match.group(1)

        if "description" in data:
            match = re.search(r"Del (\d+) av (\d+)", data["description"])
            if match:
                return match.group(1)

        return None

    def extrametadata(self, episode):
        self.output["tvshow"] = self.output["season"] is not None and self.output["episode"] is not None
        if "validFrom" in episode["item"]:

            def _fix_broken_timezone_implementation(value):
                # cx_freeze cant include .zip file for dateutil and < py37 have issues with timezones with : in it
                if "+" in value and ":" == value[-3:-2]:
                    value = value[:-3] + value[-2:]
                return value

            validfrom = episode["item"]["validFrom"]
            if "+" in validfrom:
                date = time.mktime(
                    datetime.datetime.strptime(
                        _fix_broken_timezone_implementation(episode["item"]["validFrom"].replace("Z", "")),
                        "%Y-%m-%dT%H:%M:%S%z",
                    ).timetuple(),
                )
            else:
                date = time.mktime(
                    datetime.datetime.strptime(
                        _fix_broken_timezone_implementation(episode["item"]["validFrom"].replace("Z", "")),
                        "%Y-%m-%dT%H:%M:%S",
                    ).timetuple(),
                )
            self.output["publishing_datetime"] = int(date)

        self.output["title_nice"] = episode["moreDetails"]["heading"]

        try:
            t = episode["item"]["parent"]["image"]["wide"]
        except KeyError:
            t = ""
        if isinstance(t, dict):
            url = f"https://www.svtstatic.se/image/original/default/{t['id']}/{t['changed']}?format=auto&quality=100"
            self.output["showthumbnailurl"] = url
        elif t:
            # Get the image if size/format is not specified in the URL set it to large
            url = t.format(format="large")
            self.output["showthumbnailurl"] = url
        try:
            t = episode["images"]["wide"]
        except KeyError:
            t = ""
        if isinstance(t, dict):
            url = f"https://www.svtstatic.se/image/original/default/{t['id']}/{t['changed']}?format=auto&quality=100"
            self.output["episodethumbnailurl"] = url
        elif t:
            # Get the image if size/format is not specified in the URL set it to large
            url = t.format(format="large")
            self.output["episodethumbnailurl"] = url

        if "longDescription" in episode["item"]["parent"]:
            self.output["showdescription"] = episode["item"]["parent"]["longDescription"]

        if "description" in episode:
            self.output["episodedescription"] = episode["description"]
<commit_msg>Fixes better language on error message when there's no video file found.<commit_after># ex:ts=4:sw=4:sts=4:et
# -*- tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*-
import copy
import datetime
import hashlib
import json
import logging
import re
import time
from urllib.parse import parse_qs
from urllib.parse import urljoin
from urllib.parse import urlparse

from svtplay_dl.error import ServiceError
from svtplay_dl.fetcher.dash import dashparse
from svtplay_dl.fetcher.hls import hlsparse
from svtplay_dl.service import MetadataThumbMixin
from svtplay_dl.service import Service
from svtplay_dl.subtitle import subtitle
from svtplay_dl.utils.text import filenamify

URL_VIDEO_API = "https://api.svt.se/video/"
LIVE_CHANNELS = {
    "svtbarn": "ch-barnkanalen",
    "svt1": "ch-svt1",
    "svt2": "ch-svt2",
    "svt24": "ch-svt24",
    "kunskapskanalen": "ch-kunskapskanalen",
}

class Svtplay(Service, MetadataThumbMixin):
    supported_domains = ["svtplay.se", "svt.se", "beta.svtplay.se", "svtflow.se"]
    info_search_expr = r"<script id=\"__NEXT_DATA__\" type=\"application\/json\">({.+})<\/script>"
    access = None

    def get(self):
        parse = urlparse(self.url)
        if parse.netloc in ("www.svtplay.se", "svtplay.se"):
            if parse.path[:6] != "/video" and parse.path[:6] != "/klipp" and parse.path[:8] != "/kanaler":
                yield ServiceError("This mode is not supported anymore. Need the url with the video.")
                return

        query = parse_qs(parse.query)

        if "accessService" in query:
            self.access = query["accessService"]

        urldata = self.get_urldata()

        if parse.path[:8] == "/kanaler":
            ch = LIVE_CHANNELS[parse.path[parse.path.rfind("/") + 1 :]]
            _url = urljoin(URL_VIDEO_API, ch)
            res = self.http.get(_url)
            try:
                janson = res.json()
            except json.decoder.JSONDecodeError:
                yield ServiceError(f"Can't decode api request: {res.request.url}")
                return
            videos = self._get_video(janson)
            self.config.set("live", True)
            yield from videos
            return

        match = re.search(self.info_search_expr, urldata)
        if not match:
            yield ServiceError("Can't find video info.")
            return

        janson = json.loads(match.group(1))
        video_data = None
        vid = None
        for data_entry in janson["props"]["urqlState"].values():
            if "data" in data_entry:
                entry = json.loads(data_entry["data"])
                for key, data in entry.items():
                    if key == "detailsPageByPath" and data and "moreDetails" in data:
                        video_data = data
                        vid = data["video"]["svtId"]
                        break

        if not vid:
            yield ServiceError("Can't find video id")
            return

        self.outputfilename(video_data)
        self.extrametadata(video_data)

        res = self.http.get(URL_VIDEO_API + vid)
        try:
            janson = res.json()
        except json.decoder.JSONDecodeError:
            yield ServiceError(f"Can't decode api request: {res.request.url}")
            return
        if res.status_code >= 400:
            yield ServiceError("Can't find any videos. Is it removed?")
            return
        videos = self._get_video(janson)
        yield from videos

    def _get_video(self, janson):
        if "subtitleReferences" in janson:
            for i in janson["subtitleReferences"]:
                if i["format"] == "webvtt" and "url" in i:
                    yield subtitle(copy.copy(self.config), "wrst", i["url"], "sv", output=self.output)

        if "videoReferences" in janson:
            if len(janson["videoReferences"]) == 0:
                yield ServiceError("Media doesn't have any associated videos.")
                return

            for i in janson["videoReferences"]:
                if i["url"].find(".m3u8") > 0:
                    yield from hlsparse(self.config, self.http.request("get", i["url"]), i["url"], self.output)
                elif i["url"].find(".mpd") > 0:
                    yield from dashparse(self.config, self.http.request("get", i["url"]), i["url"], output=self.output)

    def _last_chance(self):
        videos = []
        match = re.search(self.info_search_expr, self.get_urldata())
        if not match:
            logging.error("Can't find video info.")
            return videos
        janson = json.loads(match.group(1))
        video_data = None
        for data_entry in janson["props"]["urqlState"].values():
            entry = json.loads(data_entry["data"])
            for key, data in entry.items():
                if key == "startForSvtPlay":
                    video_data = data
        if not video_data:
            return videos
        for section in video_data["selections"]:
            for i in section["items"]:
                videos.append(urljoin("http://www.svtplay.se", i["item"]["urls"]["svtplay"]))
        return videos

    def _genre(self):
        videos = []
        episodes = []
        singles = []
        parse = urlparse(self.url)
        url = f"https://www.svtplay.se{parse.path}?tab=all"
        data = self.http.get(url).text
        match = re.search(self.info_search_expr, data)
        if not match:
            logging.error("Can't find video info.")
            return episodes
        janson = json.loads(match.group(1))
        video_data = None
        for data_entry in janson["props"]["urqlState"].values():
            entry = json.loads(data_entry["data"])
            for key, data in entry.items():
                if key == "categoryPage":
                    if "lazyLoadedTabs" in data:
                        video_data = data["lazyLoadedTabs"]
        if not video_data:
            return episodes
        for lazytab in video_data:
            if "selections" in lazytab:
                for section in lazytab["selections"]:
                    for i in section["items"]:
                        if i["item"]["__typename"] == "Single":
                            singles.append(urljoin("http://www.svtplay.se", i["item"]["urls"]["svtplay"]))
                        else:
                            videos.append(urljoin("http://www.svtplay.se", i["item"]["urls"]["svtplay"]))
        for i in videos:
            episodes.extend(self._all_episodes(i))
        if singles:
            episodes.extend(singles)
        return episodes

    def _all_episodes(self, url):
        parse = urlparse(url)
        tab = None
        videos = []
        if parse.query:
            query = parse_qs(parse.query)
            if "tab" in query:
                tab = query["tab"][0]

        data = self.http.get(url).text
        match = re.search(self.info_search_expr, data)
        if not match:
            logging.error("Can't find video info.")
            return videos

        janson = json.loads(match.group(1))
        video_data = None
        for data_entry in janson["props"]["urqlState"].values():
            if "data" in data_entry:
                entry = json.loads(data_entry["data"])
                # logging.info(json.dumps(entry))
                for key, data in entry.items():
                    if key == "detailsPageByPath" and data and "heading" in data:
                        video_data = data
                        break

        collections = []
        if video_data is None:
            return videos
        if video_data["item"]["parent"]["__typename"] == "Single":
            videos.append(urljoin("http://www.svtplay.se", video_data["item"]["urls"]["svtplay"]))
        for i in video_data["associatedContent"]:
            if tab:
                if tab == i["id"]:
                    collections.append(i)
            else:
                if i["id"] == "upcoming" or i["id"] == "related":
                    continue
                elif self.config.get("include_clips") and "clips" in i["id"]:
                    collections.append(i)
                elif "clips" not in i["id"]:
                    collections.append(i)

        for i in collections:
            for epi in i["items"]:
                if epi["item"]["urls"]["svtplay"] not in videos:
                    videos.append(urljoin("http://www.svtplay.se", epi["item"]["urls"]["svtplay"]))
        return videos

    def find_all_episodes(self, config):
        parse = urlparse(self._url)

        episodes = []
        if re.search("^/sista-chansen", parse.path):
            episodes = self._last_chance()
        elif re.search("^/genre", parse.path):
            episodes = self._genre()
        else:
            episodes = self._all_episodes(self.url)

        if not episodes:
            logging.error("Can't find any videos.")
        else:
            if config.get("all_last") > 0:
                return episodes[-config.get("all_last") :]
        return episodes

    def outputfilename(self, data):
        name = None
        desc = None

        name = data["moreDetails"]["heading"]
        other = data["moreDetails"]["episodeHeading"]
        vid = hashlib.sha256(data["video"]["svtId"].encode("utf-8")).hexdigest()[:7]

        if name == other:
            other = None
        elif name is None:
            name = other
            other = None
        elif other is None:
            if name != data["moreDetails"]["titleHeading"]:
                other = data["moreDetails"]["titleHeading"]

        season, episode = self.seasoninfo(data)
        if "accessibility" in data:
            if data["accessibility"] == "AudioDescribed":
                desc = "syntolkat"
            if data["accessibility"] == "SignInterpreted":
                desc = "teckentolkat"

        if not other:
            other = desc
        elif desc:
            other += f"-{desc}"

        self.output["title"] = filenamify(name)
        self.output["id"] = vid
        self.output["season"] = season
        self.output["episode"] = episode
        self.output["episodename"] = other

    def seasoninfo(self, data):
        season, episode = None, None

        season_nr = self._find_season(data)
        episode_nr = self._find_episode(data)

        if season_nr is None or episode_nr is None:
            return season, episode

        season = f"{int(season_nr):02d}"
        episode = f"{int(episode_nr):02d}"

        return season, episode

    def _find_season(self, data):
        match = re.search(r"/säsong (\d+)/", data["analyticsIdentifiers"]["viewId"])
        if match:
            return match.group(1)

        match = re.search(r"-sasong-(\d+)-", data["item"]["urls"]["svtplay"])
        if match:
            return match.group(1)

        vid = data["video"]["svtId"]
        for seasons in data["associatedContent"]:
            for i in seasons["items"]:
                if i["item"]["videoSvtId"] == vid and "positionInSeason" in i["item"]:
                    match = re.search(r"S.song (\d+)", i["item"]["positionInSeason"])
                    if match:
                        return match.group(1)
        return None

    def _find_episode(self, data):
        match = re.search(r"/avsnitt (\d+)", data["analyticsIdentifiers"]["viewId"])
        if match:
            return match.group(1)

        match = re.search(r"Avsnitt (\d+)", data["item"]["name"])
        if match:
            return match.group(1)

        vid = data["video"]["svtId"]
        for seasons in data["associatedContent"]:
            for i in seasons["items"]:
                if i["item"]["videoSvtId"] == vid and "positionInSeason" in i["item"]:
                    match = re.search(r"Avsnitt (\d+)", i["item"]["positionInSeason"])
                    if match:
                        return match.group(1)

        if "description" in data:
            match = re.search(r"Del (\d+) av (\d+)", data["description"])
            if match:
                return match.group(1)

        return None

    def extrametadata(self, episode):
        self.output["tvshow"] = self.output["season"] is not None and self.output["episode"] is not None
        if "validFrom" in episode["item"]:

            def _fix_broken_timezone_implementation(value):
                # cx_freeze cant include .zip file for dateutil and < py37 have issues with timezones with : in it
                if "+" in value and ":" == value[-3:-2]:
                    value = value[:-3] + value[-2:]
                return value

            validfrom = episode["item"]["validFrom"]
            if "+" in validfrom:
                date = time.mktime(
                    datetime.datetime.strptime(
                        _fix_broken_timezone_implementation(episode["item"]["validFrom"].replace("Z", "")),
                        "%Y-%m-%dT%H:%M:%S%z",
                    ).timetuple(),
                )
            else:
                date = time.mktime(
                    datetime.datetime.strptime(
                        _fix_broken_timezone_implementation(episode["item"]["validFrom"].replace("Z", "")),
                        "%Y-%m-%dT%H:%M:%S",
                    ).timetuple(),
                )
            self.output["publishing_datetime"] = int(date)

        self.output["title_nice"] = episode["moreDetails"]["heading"]

        try:
            t = episode["item"]["parent"]["image"]["wide"]
        except KeyError:
            t = ""
        if isinstance(t, dict):
            url = f"https://www.svtstatic.se/image/original/default/{t['id']}/{t['changed']}?format=auto&quality=100"
            self.output["showthumbnailurl"] = url
        elif t:
            # Get the image if size/format is not specified in the URL set it to large
            url = t.format(format="large")
            self.output["showthumbnailurl"] = url
        try:
            t = episode["images"]["wide"]
        except KeyError:
            t = ""
        if isinstance(t, dict):
            url = f"https://www.svtstatic.se/image/original/default/{t['id']}/{t['changed']}?format=auto&quality=100"
            self.output["episodethumbnailurl"] = url
        elif t:
            # Get the image if size/format is not specified in the URL set it to large
            url = t.format(format="large")
            self.output["episodethumbnailurl"] = url

        if "longDescription" in episode["item"]["parent"]:
            self.output["showdescription"] = episode["item"]["parent"]["longDescription"]

        if "description" in episode:
            self.output["episodedescription"] = episode["description"]
<|endoftext|>


<commit_before>'''
Rought and ready script to pull all the data from the 360Giving dkan
to a local machine.

Probably needs some work..
'''

import urllib2
import json
import httplib
import os.path
import time

# Used to debug connection info
httplib.HTTPConnection.debuglevel = 1

# Create a directory to hold the data
data_directory = 'data'
if not os.path.exists(data_directory):
    os.makedirs(data_directory)

# Fetch a list of package information from DKAN
base_url = "http://data.threesixtygiving.org/api/3/action/"
url = base_url + "package_list"
request = urllib2.Request(url)
response = urllib2.urlopen(url)
#print(response.headers.dict)
#quit()
data = json.loads(response.read())
print(data)

'''
    Loop through the packages to get the URLs of the data.
    Check if we have a copy already, if not, or if out of date
    then downlad it.

    Note we are only grabbing csv files
'''
for result in data['result']:
    # Use this for testing, jys just grab a single package
    #if result != "blf-grants-data-2004-onwards":
        #continue
    # Fetch a single package, to find the URLs of all the raw data
    url = base_url + "package_show?id=" + result
    print (result)
    response2 = urllib2.urlopen(url)
    data_package = json.loads(response2.read())

    # Find the raw data we want from the resources in the previous data.
    # We only want csvs
    for resource in data_package['result']['resources']:
        if resource['format'] == 'csv':
            print (resource['id'])
            # Now we have the resource, so go and find the URL of the data (at last!
            url = base_url + "resource_show?id=" + resource['id']
            response3 = urllib2.urlopen(url)
            #print(response3.headers.dict)
            #quit()
            data_package2 = json.loads(response3.read())
            data_url = data_package2['result']['url']
            # We have the URL!
            print(data_url)

            # Fetch the data!
            testfile = urllib2.urlopen(data_url)
            #print(testfile.headers.dict) # Get the response headers

            if testfile.headers.get('Last-Modified'):
            # Check to see if file already exists. If so see if it needs
            # grabbing again (check against last modified response)
                if os.path.isfile(data_directory + '/' + result + '.csv'):
                    file_last_modified = time.ctime(os.path.getmtime(data_directory + '/' + result + '.csv'))
                    header_last_modified = testfile.headers.get('Last-Modified')
                    print "last modified: %s" % time.ctime(os.path.getmtime(data_directory + '/' + result + '.csv'))
                    print "created: %s" % time.ctime(os.path.getctime(data_directory + '/' + result + '.csv'))
                    #print('Last Modified response: ' + testfile.headers.get('Last-Modified'))
                    if header_last_modified > file_last_modified:
                        print('Fetching file')
                        data = testfile.read()
                        with open(data_directory + '/' + result + '.csv', 'wb') as code:
                            code.write(data)
                    else:
                        print('File on disk newer than remote. Skip downloading')
                else:  # Fetch file
                    print('Fetching file')
                    data = testfile.read()
                    with open(data_directory + '/' + result + '.csv', 'wb') as code:
                        code.write(data)
            else:  # Fetch file
                print('Fetching file')
                data = testfile.read()
                with open(data_directory + '/' + result + '.csv', 'wb') as code:
                    code.write(data)
<commit_msg>Improve dkan_calls script<commit_after>'''
Rought and ready script to pull all the data from the 360Giving dkan
to a local machine.

Probably needs some work..
'''

import requests
import json
import os.path
import time

# Create a directory to hold the data
data_directory = 'data'
os.makedirs(data_directory, exist_ok=True)

# Fetch a list of package information from DKAN
base_url = "http://data.threesixtygiving.org/api/3/action/"
url = base_url + "current_package_list_with_resources"
request = requests.get(url)
data = request.json()

'''
    Loop through the packages to get the URLs of the data.
    Check if we have a copy already, if not, or if out of date
    then downlad it.

    Note we are only grabbing csv files
'''
for result in data['result']:
    # We need to make another request with the ID https://github.com/ThreeSixtyGiving/data.threesixtygiving.org/issues/11
    response3 = requests.get(base_url + 'package_show?id=' + result['id'])
    data_package = response3.json()

    # Find the raw data we want from the resources in the data.
    # We only want csvs
    for resource in data_package['result']['resources']:
        if resource['format'] == 'csv':
            print (resource['id'])
            # Now we have the resource, so go and find the URL of the data (at last!
            url = base_url + "resource_show?id=" + resource['id']
            response3 = requests.get(url)
            data_resource = response3.json()
            data_url = data_resource['result']['url']
            # We have the URL!
            print(data_url)

            # Fetch the data!
            testfile = requests.get(data_url)

            os.makedirs(os.path.join(data_directory, result['name']), exist_ok=True)
            output_name = os.path.join(data_directory, result['name'], resource['id'] + '.csv')

            with open(output_name, 'wb') as fp:
                fp.write(testfile.content)
<|endoftext|>


<commit_before>"""
Module entrypoint for core Opal views
"""
import warnings

from django.conf import settings
from django.contrib.auth.views import login
from django.http import HttpResponseNotFound
from django.shortcuts import get_object_or_404, redirect
from django.template.loader import get_template
from django.template import TemplateDoesNotExist
from django.views.generic import TemplateView, View

from opal import models
from opal.core import application, detail, episodes
from opal.core.patient_lists import PatientList, TabbedPatientListGroup
from opal.core.subrecords import (
    episode_subrecords, get_subrecord_from_api_name
)
from opal.core.views import LoginRequiredMixin, json_response
from opal.utils import camelcase_to_underscore
from opal.utils.banned_passwords import banned

warnings.simplefilter('once', DeprecationWarning)

app = application.get_app()

Synonym = models.Synonym

class PatientListTemplateView(LoginRequiredMixin, TemplateView):

    def dispatch(self, *args, **kwargs):
        try:
            self.patient_list = PatientList.get(kwargs['slug'])
        except ValueError:
            self.patient_list = None
        return super(PatientListTemplateView, self).dispatch(*args, **kwargs)

    def get_column_context(self, **kwargs):
        """
        Return the context for our columns
        """
        # we use this view to load blank tables without content for
        # the list redirect view, so if there are no kwargs, just
        # return an empty context
        if not self.patient_list:
            return []

        return self.patient_list.schema_to_dicts()

    def get_context_data(self, **kwargs):
        context = super(
            PatientListTemplateView, self
        ).get_context_data(**kwargs)
        list_slug = None
        if self.patient_list:
            list_slug = self.patient_list.get_slug()
        context['list_slug'] = list_slug
        context['patient_list'] = self.patient_list
        context['lists'] = PatientList.for_user(self.request.user)

        context['list_group'] = None
        if self.patient_list:
            group = TabbedPatientListGroup.for_list(self.patient_list)
            if group:
                if group.visible_to(self.request.user):
                    context['list_group'] = group

        context['columns'] = self.get_column_context(**kwargs)
        return context

    def get_template_names(self):
        if self.patient_list:
            return self.patient_list().get_template_names()
        return [PatientList.template_name]

class PatientDetailTemplateView(LoginRequiredMixin, TemplateView):
    template_name = 'patient_detail.html'

    def get_context_data(self, **kwargs):
        context = super(
            PatientDetailTemplateView, self
        ).get_context_data(**kwargs)

        # django likes to try and initialise classes, even when we
        # don't want it to, so vars it
        context['episode_categories'] = [
            vars(i) for i in episodes.EpisodeCategory.list()
        ]

        # We cast this to a list because it's a generator but we want to
        # consume it twice in the template
        context['detail_views'] = list(
            detail.PatientDetailView.for_user(self.request.user)
        )
        return context

# TODO: ?Remove this ?
class EpisodeDetailTemplateView(LoginRequiredMixin, TemplateView):
    def get(self, *args, **kwargs):
        self.episode = get_object_or_404(models.Episode, pk=kwargs['pk'])
        return super(EpisodeDetailTemplateView, self).get(*args, **kwargs)

    def get_template_names(self):
        names = [
            'detail/{0}.html'.format(self.episode.category_name.lower()),
            'detail/default.html'
        ]
        return names

    def get_context_data(self, **kwargs):
        context = super(
            EpisodeDetailTemplateView, self
        ).get_context_data(**kwargs)
        return context

def get_brand_name():
    warnthem = """
    The `brand_name` context variable is available via the OPAL_BRAND_NAME
setting and will be removed in Opal 0.9.0
"""
    warnings.warn(warnthem, DeprecationWarning, stacklevel=2)
    return getattr(settings, 'OPAL_BRAND_NAME', 'OPAL')

def get_settings():
    warnthem = """
    The `settings` context variable will be removed in Opal 0.9.0
    """
    warnings.warn(warnthem, DeprecationWarning, stacklevel=2)
    return settings

def get_extra_application():
    warnthem = """
    The `extra_application` context variable is available via the
OPAL_EXTRA_APPLICATION setting and will be removed in Opal 0.9.0
"""
    warnings.warn(warnthem, DeprecationWarning, stacklevel=2)
    return settings.OPAL_EXTRA_APPLICATION

class IndexView(LoginRequiredMixin, TemplateView):
    template_name = 'opal.html'

    def get_context_data(self, **kwargs):
        context = super(IndexView, self).get_context_data(**kwargs)
        context['brand_name'] = get_brand_name
        context['settings'] = get_settings
        context['extra_application'] = get_extra_application
        if hasattr(settings, 'OPAL_EXTRA_APPLICATION'):
            context['extra_application'] = get_extra_application

        return context

def check_password_reset(request, *args, **kwargs):
    """
    Check to see if the user needs to reset their password
    """
    response = login(request, *args, **kwargs)
    if response.status_code == 302:
        try:
            profile = request.user.profile
            if profile and profile.force_password_change:
                return redirect(
                    'django.contrib.auth.views.password_change'
                )
        except models.UserProfile.DoesNotExist:
            # TODO: This probably doesn't do any harm, but
            # we should really never reach this. Creation
            # of profiles shouldn't happen in a random view.
            models.UserProfile.objects.create(
                user=request.user, force_password_change=True)
            return redirect(
                'django.contrib.auth.views.password_change'
            )
    return response

"""Internal (Legacy) API View"""

class EpisodeCopyToCategoryView(LoginRequiredMixin, View):
    """
    Copy an episode to a given category, excluding tagging.
    """
    def post(self, request, pk=None, category=None, **kwargs):
        old = models.Episode.objects.get(pk=pk)
        new = models.Episode(patient=old.patient,
                             category_name=category,
                             date_of_admission=old.date_of_admission)
        new.save()

        for sub in episode_subrecords():
            if sub._is_singleton or not sub._clonable:
                continue
            for item in sub.objects.filter(episode=old):
                item.id = None
                item.episode = new
                item.save()
        serialised = new.to_dict(self.request.user)
        return json_response(serialised)

"""
Template views for OPAL
"""

class FormTemplateView(LoginRequiredMixin, TemplateView):
    """
    This view renders the form template for our field.

    These are generated for subrecords, but can also be used
    by plugins for other mdoels.
    """
    template_name = "form_base.html"

    def get_context_data(self, *args, **kwargs):
        ctx = super(FormTemplateView, self).get_context_data(*args, **kwargs)
        ctx["form_name"] = self.column.get_form_template()
        return ctx

    def dispatch(self, *a, **kw):
        """
        Set the context for what this modal is for so
        it can be accessed by all subsequent methods
        """
        self.column = get_subrecord_from_api_name(kw['model'])
        self.name = camelcase_to_underscore(self.column.__name__)
        return super(FormTemplateView, self).dispatch(*a, **kw)

class ModalTemplateView(LoginRequiredMixin, TemplateView):
    def get_template_from_model(self):
        patient_list = None

        if self.list_slug:
            patient_list = PatientList.get(self.list_slug)()
        return self.column.get_modal_template(
            patient_list=patient_list
        )

    def dispatch(self, *a, **kw):
        """
        Set the context for what this modal is for so
        it can be accessed by all subsequent methods
        """
        self.column = kw['model']
        self.list_slug = kw.get('list', None)
        self.template_name = self.get_template_from_model()
        if self.template_name is None:
            raise ValueError(
                'No modal Template available for {0}'.format(
                    self.column.__name__
                )
            )
        self.name = camelcase_to_underscore(self.column.__name__)
        return super(ModalTemplateView, self).dispatch(*a, **kw)

    def get_context_data(self, **kwargs):
        context = super(ModalTemplateView, self).get_context_data(**kwargs)
        context['name'] = self.name
        context['title'] = getattr(
            self.column, '_title', self.name.replace('_', ' ').title()
        )
        context['icon'] = getattr(self.column, '_icon', '')
        # pylint: disable=W0201
        context['single'] = self.column._is_singleton
        context["column"] = self.column

        return context

class RecordTemplateView(LoginRequiredMixin, TemplateView):
    def get_template_names(self):
        model = get_subrecord_from_api_name(self.kwargs["model"])
        template_name = model.get_display_template()
        return [template_name]

class AccountDetailTemplateView(TemplateView):
    template_name = 'accounts/account_detail.html'

class BannedView(TemplateView):
    template_name = 'accounts/banned.html'

    def get_context_data(self, *a, **k):
        data = super(BannedView, self).get_context_data(*a, **k)
        data['banned'] = banned
        return data

class HospitalNumberTemplateView(LoginRequiredMixin, TemplateView):
    template_name = 'hospital_number_modal.html'

class ReopenEpisodeTemplateView(LoginRequiredMixin, TemplateView):
    template_name = 'reopen_episode_modal.html'

class UndischargeTemplateView(LoginRequiredMixin, TemplateView):
    template_name = 'undischarge_modal.html'

class DischargeEpisodeTemplateView(LoginRequiredMixin, TemplateView):
    template_name = 'discharge_episode_modal.html'

class CopyToCategoryTemplateView(LoginRequiredMixin, TemplateView):
    template_name = 'copy_to_category.html'

class DeleteItemConfirmationView(LoginRequiredMixin, TemplateView):
    template_name = 'delete_item_confirmation_modal.html'

class RawTemplateView(LoginRequiredMixin, TemplateView):
    """
    Failover view for templates - just look for this path in Django!
    """
    def get(self, *args, **kw):
        self.template_name = kw['template_name']
        try:
            get_template(self.template_name)
        except TemplateDoesNotExist:
            return HttpResponseNotFound()
        return super(RawTemplateView, self).get(*args, **kw)
<commit_msg>De-capitalize OPAL<commit_after>"""
Module entrypoint for core Opal views
"""
import warnings

from django.conf import settings
from django.contrib.auth.views import login
from django.http import HttpResponseNotFound
from django.shortcuts import get_object_or_404, redirect
from django.template.loader import get_template
from django.template import TemplateDoesNotExist
from django.views.generic import TemplateView, View

from opal import models
from opal.core import application, detail, episodes
from opal.core.patient_lists import PatientList, TabbedPatientListGroup
from opal.core.subrecords import (
    episode_subrecords, get_subrecord_from_api_name
)
from opal.core.views import LoginRequiredMixin, json_response
from opal.utils import camelcase_to_underscore
from opal.utils.banned_passwords import banned

warnings.simplefilter('once', DeprecationWarning)

app = application.get_app()

Synonym = models.Synonym

class PatientListTemplateView(LoginRequiredMixin, TemplateView):

    def dispatch(self, *args, **kwargs):
        try:
            self.patient_list = PatientList.get(kwargs['slug'])
        except ValueError:
            self.patient_list = None
        return super(PatientListTemplateView, self).dispatch(*args, **kwargs)

    def get_column_context(self, **kwargs):
        """
        Return the context for our columns
        """
        # we use this view to load blank tables without content for
        # the list redirect view, so if there are no kwargs, just
        # return an empty context
        if not self.patient_list:
            return []

        return self.patient_list.schema_to_dicts()

    def get_context_data(self, **kwargs):
        context = super(
            PatientListTemplateView, self
        ).get_context_data(**kwargs)
        list_slug = None
        if self.patient_list:
            list_slug = self.patient_list.get_slug()
        context['list_slug'] = list_slug
        context['patient_list'] = self.patient_list
        context['lists'] = PatientList.for_user(self.request.user)

        context['list_group'] = None
        if self.patient_list:
            group = TabbedPatientListGroup.for_list(self.patient_list)
            if group:
                if group.visible_to(self.request.user):
                    context['list_group'] = group

        context['columns'] = self.get_column_context(**kwargs)
        return context

    def get_template_names(self):
        if self.patient_list:
            return self.patient_list().get_template_names()
        return [PatientList.template_name]

class PatientDetailTemplateView(LoginRequiredMixin, TemplateView):
    template_name = 'patient_detail.html'

    def get_context_data(self, **kwargs):
        context = super(
            PatientDetailTemplateView, self
        ).get_context_data(**kwargs)

        # django likes to try and initialise classes, even when we
        # don't want it to, so vars it
        context['episode_categories'] = [
            vars(i) for i in episodes.EpisodeCategory.list()
        ]

        # We cast this to a list because it's a generator but we want to
        # consume it twice in the template
        context['detail_views'] = list(
            detail.PatientDetailView.for_user(self.request.user)
        )
        return context

# TODO: ?Remove this ?
class EpisodeDetailTemplateView(LoginRequiredMixin, TemplateView):
    def get(self, *args, **kwargs):
        self.episode = get_object_or_404(models.Episode, pk=kwargs['pk'])
        return super(EpisodeDetailTemplateView, self).get(*args, **kwargs)

    def get_template_names(self):
        names = [
            'detail/{0}.html'.format(self.episode.category_name.lower()),
            'detail/default.html'
        ]
        return names

    def get_context_data(self, **kwargs):
        context = super(
            EpisodeDetailTemplateView, self
        ).get_context_data(**kwargs)
        return context

def get_brand_name():
    warnthem = """
    The `brand_name` context variable is available via the OPAL_BRAND_NAME
setting and will be removed in Opal 0.9.0
"""
    warnings.warn(warnthem, DeprecationWarning, stacklevel=2)
    return getattr(settings, 'OPAL_BRAND_NAME', 'Opal')

def get_settings():
    warnthem = """
    The `settings` context variable will be removed in Opal 0.9.0
    """
    warnings.warn(warnthem, DeprecationWarning, stacklevel=2)
    return settings

def get_extra_application():
    warnthem = """
    The `extra_application` context variable is available via the
OPAL_EXTRA_APPLICATION setting and will be removed in Opal 0.9.0
"""
    warnings.warn(warnthem, DeprecationWarning, stacklevel=2)
    return settings.OPAL_EXTRA_APPLICATION

class IndexView(LoginRequiredMixin, TemplateView):
    template_name = 'opal.html'

    def get_context_data(self, **kwargs):
        context = super(IndexView, self).get_context_data(**kwargs)
        context['brand_name'] = get_brand_name
        context['settings'] = get_settings
        context['extra_application'] = get_extra_application
        if hasattr(settings, 'OPAL_EXTRA_APPLICATION'):
            context['extra_application'] = get_extra_application

        return context

def check_password_reset(request, *args, **kwargs):
    """
    Check to see if the user needs to reset their password
    """
    response = login(request, *args, **kwargs)
    if response.status_code == 302:
        try:
            profile = request.user.profile
            if profile and profile.force_password_change:
                return redirect(
                    'django.contrib.auth.views.password_change'
                )
        except models.UserProfile.DoesNotExist:
            # TODO: This probably doesn't do any harm, but
            # we should really never reach this. Creation
            # of profiles shouldn't happen in a random view.
            models.UserProfile.objects.create(
                user=request.user, force_password_change=True)
            return redirect(
                'django.contrib.auth.views.password_change'
            )
    return response

"""Internal (Legacy) API View"""

class EpisodeCopyToCategoryView(LoginRequiredMixin, View):
    """
    Copy an episode to a given category, excluding tagging.
    """
    def post(self, request, pk=None, category=None, **kwargs):
        old = models.Episode.objects.get(pk=pk)
        new = models.Episode(patient=old.patient,
                             category_name=category,
                             date_of_admission=old.date_of_admission)
        new.save()

        for sub in episode_subrecords():
            if sub._is_singleton or not sub._clonable:
                continue
            for item in sub.objects.filter(episode=old):
                item.id = None
                item.episode = new
                item.save()
        serialised = new.to_dict(self.request.user)
        return json_response(serialised)

"""
Template views for Opal
"""

class FormTemplateView(LoginRequiredMixin, TemplateView):
    """
    This view renders the form template for our field.

    These are generated for subrecords, but can also be used
    by plugins for other mdoels.
    """
    template_name = "form_base.html"

    def get_context_data(self, *args, **kwargs):
        ctx = super(FormTemplateView, self).get_context_data(*args, **kwargs)
        ctx["form_name"] = self.column.get_form_template()
        return ctx

    def dispatch(self, *a, **kw):
        """
        Set the context for what this modal is for so
        it can be accessed by all subsequent methods
        """
        self.column = get_subrecord_from_api_name(kw['model'])
        self.name = camelcase_to_underscore(self.column.__name__)
        return super(FormTemplateView, self).dispatch(*a, **kw)

class ModalTemplateView(LoginRequiredMixin, TemplateView):
    def get_template_from_model(self):
        patient_list = None

        if self.list_slug:
            patient_list = PatientList.get(self.list_slug)()
        return self.column.get_modal_template(
            patient_list=patient_list
        )

    def dispatch(self, *a, **kw):
        """
        Set the context for what this modal is for so
        it can be accessed by all subsequent methods
        """
        self.column = kw['model']
        self.list_slug = kw.get('list', None)
        self.template_name = self.get_template_from_model()
        if self.template_name is None:
            raise ValueError(
                'No modal Template available for {0}'.format(
                    self.column.__name__
                )
            )
        self.name = camelcase_to_underscore(self.column.__name__)
        return super(ModalTemplateView, self).dispatch(*a, **kw)

    def get_context_data(self, **kwargs):
        context = super(ModalTemplateView, self).get_context_data(**kwargs)
        context['name'] = self.name
        context['title'] = getattr(
            self.column, '_title', self.name.replace('_', ' ').title()
        )
        context['icon'] = getattr(self.column, '_icon', '')
        # pylint: disable=W0201
        context['single'] = self.column._is_singleton
        context["column"] = self.column

        return context

class RecordTemplateView(LoginRequiredMixin, TemplateView):
    def get_template_names(self):
        model = get_subrecord_from_api_name(self.kwargs["model"])
        template_name = model.get_display_template()
        return [template_name]

class AccountDetailTemplateView(TemplateView):
    template_name = 'accounts/account_detail.html'

class BannedView(TemplateView):
    template_name = 'accounts/banned.html'

    def get_context_data(self, *a, **k):
        data = super(BannedView, self).get_context_data(*a, **k)
        data['banned'] = banned
        return data

class HospitalNumberTemplateView(LoginRequiredMixin, TemplateView):
    template_name = 'hospital_number_modal.html'

class ReopenEpisodeTemplateView(LoginRequiredMixin, TemplateView):
    template_name = 'reopen_episode_modal.html'

class UndischargeTemplateView(LoginRequiredMixin, TemplateView):
    template_name = 'undischarge_modal.html'

class DischargeEpisodeTemplateView(LoginRequiredMixin, TemplateView):
    template_name = 'discharge_episode_modal.html'

class CopyToCategoryTemplateView(LoginRequiredMixin, TemplateView):
    template_name = 'copy_to_category.html'

class DeleteItemConfirmationView(LoginRequiredMixin, TemplateView):
    template_name = 'delete_item_confirmation_modal.html'

class RawTemplateView(LoginRequiredMixin, TemplateView):
    """
    Failover view for templates - just look for this path in Django!
    """
    def get(self, *args, **kw):
        self.template_name = kw['template_name']
        try:
            get_template(self.template_name)
        except TemplateDoesNotExist:
            return HttpResponseNotFound()
        return super(RawTemplateView, self).get(*args, **kw)
<|endoftext|>


<commit_before># ex:ts=4:sw=4:sts=4:et
# -*- tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*-
from __future__ import absolute_import
import re
import os
import xml.etree.ElementTree as ET
import copy
import hashlib

from svtplay_dl.log import log
from svtplay_dl.service import Service, OpenGraphThumbMixin
from svtplay_dl.utils import filenamify, ensure_unicode
from svtplay_dl.utils.urllib import urlparse, urljoin
from svtplay_dl.fetcher.hds import hdsparse
from svtplay_dl.fetcher.hls import hlsparse
from svtplay_dl.fetcher.http import HTTP
from svtplay_dl.subtitle import subtitle
from svtplay_dl.error import ServiceError

class Svtplay(Service, OpenGraphThumbMixin):
    supported_domains = ['svtplay.se', 'svt.se', 'beta.svtplay.se', 'svtflow.se']

    def get(self):
        old = False

        parse = urlparse(self.url)
        if parse.netloc == "www.svtplay.se" or parse.netloc == "svtplay.se":
            if parse.path[:6] != "/video":
                yield ServiceError("This mode is not supported anymore. need the url with the video")
                return

        vid = self.find_video_id()
        if vid is None:
            yield ServiceError("Cant find video id for this video")
            return
        if re.match("^[0-9]+$", vid):
            old = True

        url = "http://www.svt.se/videoplayer-api/video/%s" % vid
        data = self.http.request("get", url)
        if data.status_code == 404:
            yield ServiceError("Can't get the json file for %s" % url)
            return

        data = data.json()
        if "live" in data:
            self.options.live = data["live"]
        if old:
            params = {"output": "json"}
            dataj = self.http.request("get", self.url, params=params).json()
        else:
            dataj = data

        if self.options.output_auto:
            self.options.service = "svtplay"
            self.options.output = self.outputfilename(dataj, self.options.output, ensure_unicode(self.get_urldata()))

        if self.exclude(self.options):
            yield ServiceError("Excluding video")
            return

        if "subtitleReferences" in data:
            for i in data["subtitleReferences"]:
                if i["format"] == "websrt":
                    yield subtitle(copy.copy(self.options), "wrst", i["url"])
        if old and dataj["video"]["subtitleReferences"]:
            try:
                suburl = dataj["video"]["subtitleReferences"][0]["url"]
            except KeyError:
                pass
            if suburl and len(suburl) > 0:
                yield subtitle(copy.copy(self.options), "wrst", suburl)

        if self.options.force_subtitle:
            return

        if len(data["videoReferences"]) == 0:
            yield ServiceError("Media doesn't have any associated videos (yet?)")
            return

        for i in data["videoReferences"]:
            if i["format"] == "hls" or i["format"] == "ios":
                streams = hlsparse(self.options, self.http.request("get", i["url"]), i["url"])
                if streams:
                    for n in list(streams.keys()):
                        yield streams[n]
            if i["format"] == "hds" or i["format"] == "flash":
                match = re.search(r"\/se\/secure\/", i["url"])
                if not match:
                    streams = hdsparse(self.options, self.http.request("get", i["url"], params={"hdcore": "3.7.0"}), i["url"])
                    if streams:
                        for n in list(streams.keys()):
                            yield streams[n]

    def find_video_id(self):
        match = re.search('data-video-id="([^"]+)"', self.get_urldata())
        if match:
            return match.group(1)
        parse = urlparse(self.url)
        match = re.search("/video/([0-9]+)/", parse.path)
        if match:
            return match.group(1)
        match = re.search("/videoEpisod-([^/]+)/", parse.path)
        if match:
            self._urldata = None
            self._url = "http://www.svtplay.se/video/%s/" % match.group(1)
            return match.group(1)
        return None

    def find_all_episodes(self, options):
        match = re.search(r'<link rel="alternate" type="application/rss\+xml" [^>]*href="([^"]+)"',
                          self.get_urldata())
        if match is None:
            match = re.findall(r'a class="play[^"]+"\s+href="(/video[^"]+)"', self.get_urldata())
            if not match:
                log.error("Couldn't retrieve episode list")
                return
            episodes = [urljoin("http://www.svtplay.se", x) for x in match]
        else:
            data = self.http.request("get", match.group(1)).content
            xml = ET.XML(data)

            episodes = [x.text for x in xml.findall(".//item/link")]
        episodes_new = []
        n = 1
        for i in episodes:
            episodes_new.append(i)
            if n == options.all_last:
                break
            n += 1
        return sorted(episodes_new)

    def outputfilename(self, data, filename, raw):
        directory = os.path.dirname(filename)
        if "statistics" in data:
            name = data["statistics"]["folderStructure"]
            if name.find(".") > 0:
                name = name[:name.find(".")]
            match = re.search("^arkiv-", name)
            if match:
                name = name.replace("arkiv-", "")
            name = filenamify(name.replace("-", "."))
            other = filenamify(data["context"]["title"])
            id = data["videoId"]
        else:
            name = data["programTitle"]
            if name.find(".") > 0:
                name = name[:name.find(".")]
            name = filenamify(name.replace(" - ", "."))
            other = filenamify(data["episodeTitle"])
            id = hashlib.sha256(data["programVersionId"]).hexdigest()[:7]

        if name == other:
            other = None
        season = self.seasoninfo(raw)
        title = name
        if season:
            title += ".%s" % season
        if other:
            title += ".%s" % other
        title += "-%s-svtplay" % id
        title = filenamify(title)
        if len(directory):
            output = os.path.join(directory, title)
        else:
            output = title
        return output

    def seasoninfo(self, data):
        match = re.search(r'play_video-area-aside__sub-title">([^<]+)<span', data)
        if match:
            line = match.group(1)
        else:
            match = re.search(r'data-title="([^"]+)"', data)
            if match:
                line = match.group(1)
            else:
                return None

        line = re.sub(" +", "", match.group(1)).replace('\n', '')
        match = re.search(r"(song(\d+)-)?Avsnitt(\d+)", line)
        if match:
            if match.group(2) is None:
                season = 1
            else:
                season = int(match.group(2))
            if season < 10:
                season = "0%s" % season
            episode = int(match.group(3))
            if episode < 10:
                episode = "0%s" % episode
            return "S%sE%s" % (season, episode)
        else:
            return None
<commit_msg>svtplay: update the data from the new vid<commit_after># ex:ts=4:sw=4:sts=4:et
# -*- tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*-
from __future__ import absolute_import
import re
import os
import xml.etree.ElementTree as ET
import copy
import hashlib

from svtplay_dl.log import log
from svtplay_dl.service import Service, OpenGraphThumbMixin
from svtplay_dl.utils import filenamify, ensure_unicode
from svtplay_dl.utils.urllib import urlparse, urljoin
from svtplay_dl.fetcher.hds import hdsparse
from svtplay_dl.fetcher.hls import hlsparse
from svtplay_dl.fetcher.http import HTTP
from svtplay_dl.subtitle import subtitle
from svtplay_dl.error import ServiceError

class Svtplay(Service, OpenGraphThumbMixin):
    supported_domains = ['svtplay.se', 'svt.se', 'beta.svtplay.se', 'svtflow.se']

    def get(self):
        old = False

        parse = urlparse(self.url)
        if parse.netloc == "www.svtplay.se" or parse.netloc == "svtplay.se":
            if parse.path[:6] != "/video":
                yield ServiceError("This mode is not supported anymore. need the url with the video")
                return

        vid = self.find_video_id()
        if vid is None:
            yield ServiceError("Cant find video id for this video")
            return
        if re.match("^[0-9]+$", vid):
            old = True

        url = "http://www.svt.se/videoplayer-api/video/%s" % vid
        data = self.http.request("get", url)
        if data.status_code == 404:
            yield ServiceError("Can't get the json file for %s" % url)
            return

        data = data.json()
        if "live" in data:
            self.options.live = data["live"]
        if old:
            params = {"output": "json"}
            dataj = self.http.request("get", self.url, params=params).json()
        else:
            dataj = data

        if self.options.output_auto:
            self.options.service = "svtplay"
            self.options.output = self.outputfilename(dataj, self.options.output, ensure_unicode(self.get_urldata()))

        if self.exclude(self.options):
            yield ServiceError("Excluding video")
            return

        if "subtitleReferences" in data:
            for i in data["subtitleReferences"]:
                if i["format"] == "websrt":
                    yield subtitle(copy.copy(self.options), "wrst", i["url"])
        if old and dataj["video"]["subtitleReferences"]:
            try:
                suburl = dataj["video"]["subtitleReferences"][0]["url"]
            except KeyError:
                pass
            if suburl and len(suburl) > 0:
                yield subtitle(copy.copy(self.options), "wrst", suburl)

        if self.options.force_subtitle:
            return

        if len(data["videoReferences"]) == 0:
            yield ServiceError("Media doesn't have any associated videos (yet?)")
            return

        for i in data["videoReferences"]:
            if i["format"] == "hls" or i["format"] == "ios":
                streams = hlsparse(self.options, self.http.request("get", i["url"]), i["url"])
                if streams:
                    for n in list(streams.keys()):
                        yield streams[n]
            if i["format"] == "hds" or i["format"] == "flash":
                match = re.search(r"\/se\/secure\/", i["url"])
                if not match:
                    streams = hdsparse(self.options, self.http.request("get", i["url"], params={"hdcore": "3.7.0"}), i["url"])
                    if streams:
                        for n in list(streams.keys()):
                            yield streams[n]

    def find_video_id(self):
        match = re.search('data-video-id="([^"]+)"', self.get_urldata())
        if match:
            return match.group(1)
        parse = urlparse(self.url)
        match = re.search("/video/([0-9]+)/", parse.path)
        if match:
            return match.group(1)
        match = re.search("/videoEpisod-([^/]+)/", parse.path)
        if match:
            self._urldata = None
            self._url = "http://www.svtplay.se/video/%s/" % match.group(1)
            self.get_urldata()
            return self.find_video_id()
        return None

    def find_all_episodes(self, options):
        match = re.search(r'<link rel="alternate" type="application/rss\+xml" [^>]*href="([^"]+)"',
                          self.get_urldata())
        if match is None:
            match = re.findall(r'a class="play[^"]+"\s+href="(/video[^"]+)"', self.get_urldata())
            if not match:
                log.error("Couldn't retrieve episode list")
                return
            episodes = [urljoin("http://www.svtplay.se", x) for x in match]
        else:
            data = self.http.request("get", match.group(1)).content
            xml = ET.XML(data)

            episodes = [x.text for x in xml.findall(".//item/link")]
        episodes_new = []
        n = 1
        for i in episodes:
            episodes_new.append(i)
            if n == options.all_last:
                break
            n += 1
        return sorted(episodes_new)

    def outputfilename(self, data, filename, raw):
        directory = os.path.dirname(filename)
        if "statistics" in data:
            name = data["statistics"]["folderStructure"]
            if name.find(".") > 0:
                name = name[:name.find(".")]
            match = re.search("^arkiv-", name)
            if match:
                name = name.replace("arkiv-", "")
            name = filenamify(name.replace("-", "."))
            other = filenamify(data["context"]["title"])
            id = data["videoId"]
        else:
            name = data["programTitle"]
            if name.find(".") > 0:
                name = name[:name.find(".")]
            name = filenamify(name.replace(" - ", "."))
            other = filenamify(data["episodeTitle"])
            id = hashlib.sha256(data["programVersionId"]).hexdigest()[:7]

        if name == other:
            other = None
        season = self.seasoninfo(raw)
        title = name
        if season:
            title += ".%s" % season
        if other:
            title += ".%s" % other
        title += "-%s-svtplay" % id
        title = filenamify(title)
        if len(directory):
            output = os.path.join(directory, title)
        else:
            output = title
        return output

    def seasoninfo(self, data):
        match = re.search(r'play_video-area-aside__sub-title">([^<]+)<span', data)
        if match:
            line = match.group(1)
        else:
            match = re.search(r'data-title="([^"]+)"', data)
            if match:
                line = match.group(1)
            else:
                return None

        line = re.sub(" +", "", match.group(1)).replace('\n', '')
        match = re.search(r"(song(\d+)-)?Avsnitt(\d+)", line)
        if match:
            if match.group(2) is None:
                season = 1
            else:
                season = int(match.group(2))
            if season < 10:
                season = "0%s" % season
            episode = int(match.group(3))
            if episode < 10:
                episode = "0%s" % episode
            return "S%sE%s" % (season, episode)
        else:
            return None
<|endoftext|>


<commit_before>from galpy.orbit import Orbit
from galpy.potential import MWPotential2014 as mp
import matplotlib.pyplot as plt
import numpy as np

# GALPY STORES ORBITAL INFO IN GALACTOCENTRIC CYLINDRICAL COORDINATES
# [R, vR, vT, z, vZ, phi]

def galpy_coords_to_xyzuvw(data, ro=8., vo=220.):
    """
    Converts orbits from galactocentric and takes them to XYZUVW

    Data should be raw galpy data (i.e. output from o.getOrbit()).
    The XYZUVW will be a corotating reference frame centred on the LSR
    as defined by the Schoenrich solar motion of
    XYZUVW = 0, 0, 25pc, 11.1 km/s, 12.24 km/s, 7.25 km/s

    Parameters
    ----------
    data : [ntimes, 6] float array
        output from o.getOrbit. Data is encoded as:
        [R, vR, vT, z, vz, phi]
        R : galactic radial distance /ro
        vR : galactic radial velocity /vo
        vT : circular velocity /vo
        z  : vertical distance from plane / ro
        vz : vertical velocity / vo
        phi : angle about the galaxy
    ro : float
        a conversion factor that takes units from galpy units to
        physical units. If left as default, output will be in kpc
    vo : float
        a conversion factor that takes units form galpy units to
        physical units. If left as default, output will be in km/s
        This is also the circular velocity of a circular orbit with X,Y
        equal to that of the sun.
    """
    return None

#def demo_lsr_and_sun_calc():
if __name__ == '__main__':
    max_age = 1 #Gyr
    ntimes = 100
    ts= np.linspace(0,max_age,ntimes)

    # INITIALISING SUN COORDINATES AND ORBIT
    #deg, deg, kpc,  mas/yr, mas/yr, km/s
    ra,   dec, dist, mu_ra,  mu_dec, vlos = 0., 0., 0., 0., 0., 0.
    solar_coords = [ra, dec, dist, mu_ra, mu_dec, vlos]
    sun = Orbit(vxvv=solar_coords, radec=True, solarmotion='schoenrich') # should just be the sun's orbit
    sun.integrate(ts,mp,method='odeint')

    # get the orbit (not sure what coords...)
    sun_data = sun.getOrbit()

    # plots the sun's motion with respect to LSR(?)
    sunR = 8 * sun_data[:,0]
    sunphi = sun_data[:,5]
    sunX = sunR * np.cos(sunphi)
    sunY = sunR * np.sin(sunphi)
    sunZ = 8 * sun_data[:,3]
    plt.clf()
    plt.plot(sunX, sunY)
    plt.savefig('temp_plots/sunXY.png')

    plt.clf()
    plt.plot(sunX, sunZ)
    plt.savefig('temp_plots/sunXZ.png')

    # plot the XY of the sun's motion using galpy's plot function
    plt.clf()
    sun.plot(d1='x', d2='y')
    plt.savefig('temp_plots/galpy_sunXY.png')

    sun.plot(d1='x', d2='z')
    plt.savefig('temp_plots/galpy_sunXZ.png')

    #                                                        kpc, km/s
    # INITIALISING THE LSR (at XYZUVW (w.r.t sun) of [0,0,-0.025,0,220,0]
    R, vR, vT, z, vz, phi = 1., 0., 1., 0., 0./220, 0.
    LSR_coords = [R, vR, vT, z, vz, phi]
    lsr = Orbit(vxvv=LSR_coords, solarmotion='schoenrich', vo=220, ro=8)
    lsr.integrate(ts, mp, method='odeint')

    plt.clf()
    # SHOULD PLOT A PERFECT CIRCLE
    lsr.plot(d1='x', d2='y')
    plt.savefig('temp_plots/galpy_lsrXY.png')

    plt.clf()
    lsr.plot(d1='x', d2='z')
    plt.savefig('temp_plots/galpy_lsrXZ.png')

    lsr_data = lsr.getOrbit()
    lsrR = 8 * lsr_data[:,0]
    lsrphi = lsr_data[:,5]

    lsrX = lsrR * np.cos(lsrphi)
    lsrY = lsrR * np.sin(lsrphi)
    lsrZ = 8 * lsr_data[:,3]

    plt.clf()
    plt.plot(lsrX, lsrY)
    plt.savefig('temp_plots/lsrXY.png')
    plt.clf()
    plt.plot(lsrX, lsrZ)
    plt.savefig('temp_plots/lsrXZ.png')

    # plotting both sun and lsr
    plt.clf()
    plt.plot(lsrX, lsrY)
    plt.plot(sunX, sunY)
    plt.savefig('temp_plots/combXY.png')
    plt.clf()
    plt.plot(lsrX, lsrZ)
    plt.plot(sunX, sunZ)
    plt.savefig('temp_plots/combXZ.png')

    # Finding sun's path w.r.t the LSR in non-corotating frame
    relsunX = sunX - lsrX
    relsunY = sunY - lsrY
    relsunZ = sunZ - lsrZ

    plt.clf()
    plt.plot(relsunX, relsunY)
    plt.savefig('temp_plots/relsunXY.png')
    plt.clf()
    plt.plot(relsunX, relsunZ)
    plt.savefig('temp_plots/relsunXZ.png')

    # Getting sun's path w.r.t the LSR in cortating frame
    sun_rel_data = sun_data - lsr_data
    sun_relR = 8 * sun_rel_data[:,0]
    sun_relphi = sun_rel_data[:,5]

    sun_relX = sun_relR * np.cos(sun_relphi)
    sun_relY = sun_relR * np.sin(sun_relphi)
    sun_relZ = 8 * sun_rel_data[:,3]

    plt.clf()
    plt.plot(sun_relX, sun_relY)
    plt.savefig('temp_plots/sun_relXY.png')
    plt.clf()
    plt.plot(sun_relX, sun_relZ)
    plt.savefig('temp_plots/sun_relXZ.png')

    # Try and plot LSR and sun in 3D for comparison with
    # relative plot
    plt.clf()
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    #theta = np.linspace(-4 * np.pi, 4 * np.pi, 100)
    #z = np.linspace(-2, 2, 100)
    #r = z ** 2 + 1
    #x = r * np.sin(theta)
    #y = r * np.cos(theta)
    ax.plot(sunX, sunY, sunZ, label='sun')
    ax.plot(lsrX, lsrY, lsrZ, label='lsr')
    ax.legend()
    plt.savefig('temp_plots/3D_sun_lsr.png')
    #plt.show()

#if __name__ == '__main__':
def lsr_nonsense():
    max_age = 4. #Gyr
    ntimes = 100
    ts= np.linspace(0,max_age,ntimes)

    # demo a star (with vT=220, vR=0, vZ=0, z=0, phi=0.1 pi) staying
    # fixed in our coordinate frame
    R, vR, vT, z, vz, phi = 1., 0., 1., 0., 0./220, 0.
    LSR_coords = [R, vR, vT, z, vz, phi]
    lsr = Orbit(vxvv=LSR_coords, solarmotion='schoenrich', vo=220, ro=8)
    lsr.integrate(ts, mp, method='odeint')

    lsr_data = lsr.getOrbit()
    lsrR = 8 * lsr_data[:,0]
    lsrphi = lsr_data[:,5]

    lsrX = lsrR * np.cos(lsrphi)
    lsrY = lsrR * np.sin(lsrphi)
    lsrZ = 8 * lsr_data[:,3]

    R, vR, vT, z, vz, phi = 1., 0., 1., 0., 0./220, 0.25*np.pi
    rot_lsr_coords = [R, vR, vT, z, vz, phi]
    rot_lsr = Orbit(vxvv=rot_lsr_coords, solarmotion='schoenrich', vo=220, ro=8)
    rot_lsr.integrate(ts, mp, method='odeint')

    rot_lsr_data = rot_lsr.getOrbit()
    rot_lsrR = 8 * rot_lsr_data[:,0]
    rot_lsrphi = rot_lsr_data[:,5]

    rot_lsrX = rot_lsrR * np.cos(rot_lsrphi)
    rot_lsrY = rot_lsrR * np.sin(rot_lsrphi)
    rot_lsrZ = 8 * rot_lsr_data[:,3]

    plt.clf()
    plt.plot(lsrX, lsrY + 1)
    plt.plot(rot_lsrX, rot_lsrY)
    plt.savefig('temp_plots/rotXY.png')
    plt.clf()
    plt.plot(lsrX, lsrZ)
    plt.plot(rot_lsrX, rot_lsrZ)
    plt.savefig('temp_plots/rotXZ.png')

    # putting into corotating cartesian system centred on LSR
    rel_data = rot_lsr_data - lsr_data

<commit_msg>i dunno... changes and shit<commit_after>from galpy.orbit import Orbit
from galpy.potential import MWPotential2014 as mp
from galpy.util import bovy_conversion
import matplotlib.pyplot as plt
import numpy as np
import astropy.units as u

from IPython.core.debugger import Tracer

# GALPY STORES ORBITAL INFO IN GALACTOCENTRIC CYLINDRICAL COORDINATES
# [R, vR, vT, z, vZ, phi]

def observed_to_xyzuvw(obs, ts, lsr_orbit=None):
    """
    Convert six-parameter astrometric solution to XYZUVW orbit.

    Parameters
    ----------
    obs :   [RA (deg), DEC (deg), pi (mas),
             mu_ra (mas/yr), mu_dec (mas/yr), vlos (km/s)]
        Current kinematics
    ts : [ntimes] array
        times (in Gyr) to traceback to

    lsr_orbit : Orbit
        the orbit of the local standard of rest for comparison, if None can
        calculate on the fly

    XYZUVW : [ntimes, 6] array
        The space position and velocities of the star in a co-rotating frame
        centred on the LSR
    """

    Tracer()()

def galpy_coords_to_xyzuvw(data, ro=8., vo=220., rc=True):
    """
    Converts orbits from galactocentric and takes them to XYZUVW

    Data should be raw galpy data (i.e. output from o.getOrbit()).
    The XYZUVW will be a corotating reference frame centred on the LSR
    as defined by the Schoenrich solar motion of
    XYZUVW = 0, 0, 25pc, 11.1 km/s, 12.24 km/s, 7.25 km/s

    Parameters
    ----------
    data : [ntimes, 6] float array
        output from o.getOrbit. Data is encoded as:
        [R, vR, vT, z, vz, phi]
        R : galactic radial distance /ro
        vR : galactic radial velocity /vo
        vT : circular velocity /vo
        z  : vertical distance from plane / ro
        vz : vertical velocity / vo
        phi : angle about the galaxy
    ro : float
        a conversion factor that takes units from galpy units to
        physical units. If left as default, output will be in kpc
    vo : float
        a conversion factor that takes units form galpy units to
        physical units. If left as default, output will be in km/s
        This is also the circular velocity of a circular orbit with X,Y
        equal to that of the sun.
    rc : boolean
        whether to calculate XYZUVW in a right handed coordinate system
        (X, U positive towards galactic centre)
    """
    R, vR, vT, z, vz, phi = data.T
    phi = data[:, 5]
    X = ro * R * np.cos(phi)
    Y = ro * R * np.sin(phi)
    Z = ro * z

    U = vo * (vT * np.cos(phi) + vR * np.sin(phi))
    V = vo * (-vR * np.cos(phi) + vT * np.sin(phi))
    W = vo * vz
    XYZUVW = np.vstack((X,Y,Z,U,V,W)).T
    return XYZUVW

def demo_lsr_and_sun_cal():
    """
    Litte demo showing how one would calculate the orbit of the LSR and sun
    :return:
    """
    perimeter = 2 * np.pi * 8 * u.kpc
    velocity  = 220 * u.km/ u.s
    # for reference, LSR (at 8 kpc, with V = 220 km/s) should take this long
    # to complete one orbit
    orbit_time = (perimeter / velocity).to("Gyr")

    max_age = 100 * orbit_time / bovy_conversion.time_in_Gyr(220., 8.) # Gyr
    ntimes = 10000
    ts = np.linspace(0,max_age,ntimes)

    # INITIALISING SUN COORDINATES AND ORBIT
    #deg, deg, kpc,  mas/yr, mas/yr, km/s
    ra, dec, dist, mu_ra,  mu_dec, vlos = 0., 0., 0., 0., 0., 0.
    solar_coords = [ra, dec, dist, mu_ra, mu_dec, vlos]
    sun = Orbit(vxvv=solar_coords, radec=True, solarmotion='schoenrich') # should just be the sun's orbit
    sun.integrate(ts,mp,method='odeint')

    # get the orbit [R, vR, vT, z, vz, phi] (pos scaled by ro, vel scaled by vo)
    sun_data = sun.getOrbit()

    # plots the sun's motion with respect to Galactic Centre
    sunR = 8 * sun_data[:,0]
    sunphi = sun_data[:,5]
    sunX = sunR * np.cos(sunphi)
    sunY = sunR * np.sin(sunphi)
    sunZ = 8 * sun_data[:,3]
    plt.clf()
    plt.plot(sunX, sunY)
    plt.savefig('temp_plots/sunXY.png')

    plt.clf()
    plt.plot(sunX, sunZ)
    plt.savefig('temp_plots/sunXZ.png')

    # plot the XY of the sun's motion using galpy's plot function (w.r.t GC)
    plt.clf()
    sun.plot(d1='x', d2='y')
    plt.savefig('temp_plots/galpy_sunXY.png')

    sun.plot(d1='x', d2='z')
    plt.savefig('temp_plots/galpy_sunXZ.png')

    plt.clf()
    sun.plot(d1='R', d2='z')
    plt.savefig('temp_plots/galpy_sunRZ.png')

    #                                                        kpc, km/s
    # INITIALISING THE LSR (at XYZUVW (w.r.t sun) of [0,0,-0.025,0,220,0]
    R, vR, vT, z, vz, phi = 1., 0., 1., 0., 0., 0. # <--- Galpy units
    LSR_coords = [R, vR, vT, z, vz, phi]
    lsr = Orbit(vxvv=LSR_coords, solarmotion='schoenrich', vo=220, ro=8)
    lsr.integrate(ts, mp, method='odeint')

    # plots a perfect circle
    plt.clf()
    lsr.plot(d1='x', d2='y')
    plt.savefig('temp_plots/galpy_lsrXY.png')

    plt.clf()
    lsr.plot(d1='x', d2='z')
    plt.savefig('temp_plots/galpy_lsrXZ.png')

    # Manually reconstructing orbit
    lsr_data = lsr.getOrbit()
    lsrR = 8 * lsr_data[:,0]
    lsrphi = lsr_data[:,5]

    lsrX = lsrR * np.cos(lsrphi)
    lsrY = lsrR * np.sin(lsrphi)
    lsrZ = 8 * lsr_data[:,3]

    plt.clf()
    plt.plot(lsrX, lsrY)
    plt.savefig('temp_plots/lsrXY.png')
    plt.clf()
    plt.plot(lsrX, lsrZ)
    plt.savefig('temp_plots/lsrXZ.png')

    # plotting both sun and lsr
    plt.clf()
    plt.plot(lsrX, lsrY)
    plt.plot(sunX, sunY)
    plt.savefig('temp_plots/combXY.png')
    plt.clf()
    plt.plot(lsrX, lsrZ)
    plt.plot(sunX, sunZ)
    plt.savefig('temp_plots/combXZ.png')

    # Finding sun's path w.r.t the LSR in non-corotating frame
    relsunX = sunX - lsrX
    relsunY = sunY - lsrY
    relsunZ = sunZ - lsrZ

    plt.clf()
    plt.plot(relsunX, relsunY)
    plt.savefig('temp_plots/relsunXY.png')
    plt.clf()
    plt.plot(relsunX, relsunZ)
    plt.savefig('temp_plots/relsunXZ.png')

    # Getting sun's path w.r.t the LSR in cortating frame
    sun_rel_data = sun_data - lsr_data
    sun_relR = 8 * sun_rel_data[:,0]
    sun_relphi = sun_rel_data[:,5]

    sun_relX = sun_relR * np.cos(sun_relphi)
    sun_relY = sun_relR * np.sin(sun_relphi)
    sun_relZ = 8 * sun_rel_data[:,3]

    plt.clf()
    plt.plot(sun_relX, sun_relY)
    plt.savefig('temp_plots/sun_relXY.png')
    plt.clf()
    plt.plot(sun_relX, sun_relZ)
    plt.savefig('temp_plots/sun_relXZ.png')

    # Try and plot LSR and sun in 3D for comparison with
    # relative plot
    plt.clf()
    fig = plt.figure()
    ax = fig.gca(projection='3d')
    #theta = np.linspace(-4 * np.pi, 4 * np.pi, 100)
    #z = np.linspace(-2, 2, 100)
    #r = z ** 2 + 1
    #x = r * np.sin(theta)
    #y = r * np.cos(theta)
    ax.plot(sunX, sunY, sunZ, label='sun')
    ax.plot(lsrX, lsrY, lsrZ, label='lsr')
    ax.legend()
    plt.savefig('temp_plots/3D_sun_lsr.png')
    plt.show()
    galpy_coords_to_xyzuvw(lsr_data)
    print("Max age is {} and max phi is {}... does this make sense?".\
        format(max_age, np.max(lsr_data[:,5])))
    print("Max age is {} and max phi is {}... does this make sense?". \
        format(max_age, np.max(sun_data[:,5])))

if __name__ == '__main__':
#def lsr_nonsense():
    demo_lsr_and_sun_cal()

    RO = 8.
    VO = 220.
    BOVY_TIME_CONVERSION = bovy_conversion.time_in_Gyr(VO, RO)

    perimeter = 2 * np.pi * 8 * u.kpc
    velocity = 220 * u.km / u.s
    # for reference, LSR (at 8 kpc, with V = 220 km/s) should take this long
    # to complete one orbit
    orbit_time = (perimeter / velocity).to("Gyr")

    max_age = orbit_time.value / BOVY_TIME_CONVERSION
    ntimes = 100
    ts = np.linspace(0, max_age, ntimes)

    # demo a star (with vT=220, vR=0, vZ=0, z=0, phi=0.1 pi) staying
    # fixed in our coordinate frame
    R, vR, vT, z, vz, phi = 1., 0., 1., 0., 0., 0.
    LSR_coords = [R, vR, vT, z, vz, phi]
    lsr = Orbit(vxvv=LSR_coords, solarmotion='schoenrich', vo=220, ro=8)
    lsr.integrate(ts, mp, method='odeint')

    lsr_data = lsr.getOrbit()
    lsrR = RO * lsr_data[:,0]
    lsrphi = lsr_data[:,5]

    lsrX = lsrR * np.cos(lsrphi)
    lsrY = lsrR * np.sin(lsrphi)
    lsrZ = RO * lsr_data[:,3]

    R, vR, vT, z, vz, phi = 1., 0., 1., 0., 0., 0.25*np.pi
    rot_lsr_coords = [R, vR, vT, z, vz, phi]
    rot_lsr = Orbit(vxvv=rot_lsr_coords, solarmotion='schoenrich', vo=220, ro=8)
    rot_lsr.integrate(ts, mp, method='odeint')

    rot_lsr_data = rot_lsr.getOrbit()
    rot_lsrR = 8 * rot_lsr_data[:,0]
    rot_lsrphi = rot_lsr_data[:,5]

    rot_lsrX = rot_lsrR * np.cos(rot_lsrphi)
    rot_lsrY = rot_lsrR * np.sin(rot_lsrphi)
    rot_lsrZ = 8 * rot_lsr_data[:,3]

    plt.clf()
    plt.plot(lsrX, lsrY + 1)
    plt.plot(rot_lsrX, rot_lsrY)
    plt.savefig('temp_plots/rotXY.png')
    plt.clf()
    plt.plot(lsrX, lsrZ)
    plt.plot(rot_lsrX, rot_lsrZ)
    plt.savefig('temp_plots/rotXZ.png')

    # putting into corotating cartesian system centred on LSR
    rel_data = rot_lsr_data - lsr_data
    XYZUVW_rot = galpy_coords_to_xyzuvw(rel_data)

    #observed_to_xyzuvw(0, np.linspace(0,orbit_time.value, 100))
    galpy_ts = np.linspace(0, orbit_time.value, 100) / BOVY_TIME_CONVERSION
<|endoftext|>


<commit_before>from __future__ import division

import pickle
import os
import sys
import numpy as np
from ..constants import molwt
from ..RCPs.rcp45 import Emissions as r45e

def Stevens(emissions, stevens_params=np.array([0.001875, 0.634, 60.]),
    ref_isSO2=True):
    """Calculates aerosol forcing based on Stevens (2015) that relates sulphate
    aerosol forcing to SOx emissions in a logarithmic fashion.

    Input:
        emissions:   anthropogenic emissions database
    Keywords:
        stevens_params: 3 element array
            0. natural emissions of SOx in Mt/yr
            1. scaling parameter for ERFari (alpha)
            2. scaling parameter for ERFaci (beta)
        ref_isSO2:   True if E_SOx_nat is in units of SO2 rather than S.
    Output:
        F:           aerosol effective radiative forcing
    """

    E_SOx_nat, alpha, beta = stevens_params

    factor = 1
    if ref_isSO2:
        factor = molwt.SO2/molwt.S
    em_SOx = emissions[:,5] * factor

    ERFari = -alpha * em_SOx
    ERFaci = -beta * np.log(em_SOx/E_SOx_nat + 1)
    F = ERFari + ERFaci
    return F

def aerocom_direct(emissions,
        beta = np.array(
          [-6.2227e-3, 0.0, -3.8392e-4, -1.16551e-3, 1.601537e-2, -1.45339e-3,
           -1.55605e-3])
    ):

    """Calculates direct aerosol forcing based on linear relationships between
    emissions and forcing in Aerocom models.

    Reference: Myhre et al., 2013: https://www.atmos-chem-phys.net/13/1853/2013

    If inputs from an RCPs SCEN file are used, the units will be correct.

    Inputs:
        emissions: (nt x 40) emissions array
    Keywords:
        beta: 7-element array of forcing efficiencies in W m-2 (Mt yr-1)-1 for
            SOx, CO, NMVOC, NOx, BC, OC, NH3 (in that order)
    Outputs:
        Forcing time series
    """

    em_SOx, em_CO, em_NMVOC, em_NOx, em_BC, em_OC, em_NH3 = \
        emissions[:,[5, 6, 7, 8, 9, 10, 11]].T

    F_SOx    = beta[0] * em_SOx
    F_CO     = beta[1] * em_CO
    F_NMVOC  = beta[2] * em_NMVOC
    F_NOx    = beta[3] * em_NOx
    F_BC     = beta[4] * em_BC
    F_OC     = beta[5] * em_OC
    F_NH3    = beta[6] * em_NH3

    F = F_SOx+F_CO+F_NMVOC+F_NOx+F_BC+F_OC+F_NH3

    return F

def ghan_indirect(emissions, fix_pre1850_RCP=True, scale_AR5=False,
    ghan_params=np.array([-1.95011431, 0.01107147, 0.01387492])):
    """Estimates the aerosol indirect effect based on the simple model in
    Ghan et al., (2013), doi:10.1002/jgrd.50567.

    This function is just an emulator - a full implementation in Python of the
    Ghan routine (originally coded in Fortran) exists, but will require
    optimisation before it can be used in FAIR. I hope to make the full version
    available in a future version.

    A 500-member Latin Hypercube sample of emissions of SOx, NMVOC, BC and OC
    was prepared offline and run through the Ghan simple model and a functional
    relationship fit to the output. SOA aerosol (based on NMVOC emissions) is
    sometimes unreliable and does not exert a strong dependence on the ERF, and
    OC+BC is parameterised as primary organic matter, so the resulting output
    is a function of SOx and (BC+OC) emissions.

    Inputs:
        emissions: (nt x 40) numpy emissions array
    Keywords:
        fix_pre1850_RCP: Use different relationship for 1750/65 to 1850 based
                         on anthropogenic emissions from Skeie et al (2011)
                         for 1750 (atmos-chem-phys.net/11/11827/2011)
        scale_AR5:       If True, scale the forcing output so that the best
                         estimate forcing in 2011 is -0.45 W/m2 based on 2011
                         emissions from the RCPs. The Ghan emulator is built on
                         results from the CAM5 GCM. As reported in AR5 WG1 Ch7,
                         GCMs tend to overestimate forcing from aerosol-cloud
                         interactions.
        ghan_params:     3-element numpy array
                         0: scale factor
                         1: sensitivity to SOx emissions
                         2: sensitivity to BC+OC emissions
    Outputs:
        Forcing timeseries
    """

    year, em_SOx, em_BC, em_OC = emissions[:,[0, 5, 9, 10]].T

    def _ERFaci(em,
        ghan_params=np.array([-1.95011431, 0.01107147, 0.01387492])):
        scale = ghan_params[0]
        b_SOx = ghan_params[1]
        b_POM = ghan_params[2]
        return scale*np.log(1+b_SOx*em[0]+b_POM*em[1])

    # PI forcing was not zero as there were some emissions. Use estimates
    # from Skeie et al, 2011 for 1750 forcing.
    E_1765 = np.array([1.0, 11.2])
    nt = len(year)
    F_pd = np.zeros(nt)
    for i in range(nt):
        if year[i]>=1850 or fix_pre1850_RCP==False:
            F_pd[i] = _ERFaci([em_SOx[i], em_BC[i]+em_OC[i]],
                              ghan_params=ghan_params)
        else:
            # linearly interpolate between 1765 and 1850
            E_1850 = np.array([r45e.sox[85], r45e.bc[85]+r45e.oc[85]])
            F_pd[i] = _ERFaci((year[i]-1765)/85.*E_1850 +
                              (1850-year[i])/85.*E_1765,
                              ghan_params=ghan_params)

    # 1765 emissions = zero forcing
    F_1765 = -0.3002836449793625
    F_2011 = -1.5236182344467388

    # are we rescaling to AR5 best estimate with the default parameters?
    if scale_AR5:
        scale=-0.45/(F_2011-F_1765)
    else:
        scale=1.0

    return (F_pd - F_1765) * scale
<commit_msg>Remove unused imports in aerosols.py<commit_after>from __future__ import division

import numpy as np
from ..constants import molwt
from ..RCPs.rcp45 import Emissions as r45e

def Stevens(emissions, stevens_params=np.array([0.001875, 0.634, 60.]),
    ref_isSO2=True):
    """Calculates aerosol forcing based on Stevens (2015) that relates sulphate
    aerosol forcing to SOx emissions in a logarithmic fashion.

    Input:
        emissions:   anthropogenic emissions database
    Keywords:
        stevens_params: 3 element array
            0. natural emissions of SOx in Mt/yr
            1. scaling parameter for ERFari (alpha)
            2. scaling parameter for ERFaci (beta)
        ref_isSO2:   True if E_SOx_nat is in units of SO2 rather than S.
    Output:
        F:           aerosol effective radiative forcing
    """

    E_SOx_nat, alpha, beta = stevens_params

    factor = 1
    if ref_isSO2:
        factor = molwt.SO2/molwt.S
    em_SOx = emissions[:,5] * factor

    ERFari = -alpha * em_SOx
    ERFaci = -beta * np.log(em_SOx/E_SOx_nat + 1)
    F = ERFari + ERFaci
    return F

def aerocom_direct(emissions,
        beta = np.array(
          [-6.2227e-3, 0.0, -3.8392e-4, -1.16551e-3, 1.601537e-2, -1.45339e-3,
           -1.55605e-3])
    ):

    """Calculates direct aerosol forcing based on linear relationships between
    emissions and forcing in Aerocom models.

    Reference: Myhre et al., 2013: https://www.atmos-chem-phys.net/13/1853/2013

    If inputs from an RCPs SCEN file are used, the units will be correct.

    Inputs:
        emissions: (nt x 40) emissions array
    Keywords:
        beta: 7-element array of forcing efficiencies in W m-2 (Mt yr-1)-1 for
            SOx, CO, NMVOC, NOx, BC, OC, NH3 (in that order)
    Outputs:
        Forcing time series
    """

    em_SOx, em_CO, em_NMVOC, em_NOx, em_BC, em_OC, em_NH3 = \
        emissions[:,[5, 6, 7, 8, 9, 10, 11]].T

    F_SOx    = beta[0] * em_SOx
    F_CO     = beta[1] * em_CO
    F_NMVOC  = beta[2] * em_NMVOC
    F_NOx    = beta[3] * em_NOx
    F_BC     = beta[4] * em_BC
    F_OC     = beta[5] * em_OC
    F_NH3    = beta[6] * em_NH3

    F = F_SOx+F_CO+F_NMVOC+F_NOx+F_BC+F_OC+F_NH3

    return F

def ghan_indirect(emissions, fix_pre1850_RCP=True, scale_AR5=False,
    ghan_params=np.array([-1.95011431, 0.01107147, 0.01387492])):
    """Estimates the aerosol indirect effect based on the simple model in
    Ghan et al., (2013), doi:10.1002/jgrd.50567.

    This function is just an emulator - a full implementation in Python of the
    Ghan routine (originally coded in Fortran) exists, but will require
    optimisation before it can be used in FAIR. I hope to make the full version
    available in a future version.

    A 500-member Latin Hypercube sample of emissions of SOx, NMVOC, BC and OC
    was prepared offline and run through the Ghan simple model and a functional
    relationship fit to the output. SOA aerosol (based on NMVOC emissions) is
    sometimes unreliable and does not exert a strong dependence on the ERF, and
    OC+BC is parameterised as primary organic matter, so the resulting output
    is a function of SOx and (BC+OC) emissions.

    Inputs:
        emissions: (nt x 40) numpy emissions array
    Keywords:
        fix_pre1850_RCP: Use different relationship for 1750/65 to 1850 based
                         on anthropogenic emissions from Skeie et al (2011)
                         for 1750 (atmos-chem-phys.net/11/11827/2011)
        scale_AR5:       If True, scale the forcing output so that the best
                         estimate forcing in 2011 is -0.45 W/m2 based on 2011
                         emissions from the RCPs. The Ghan emulator is built on
                         results from the CAM5 GCM. As reported in AR5 WG1 Ch7,
                         GCMs tend to overestimate forcing from aerosol-cloud
                         interactions.
        ghan_params:     3-element numpy array
                         0: scale factor
                         1: sensitivity to SOx emissions
                         2: sensitivity to BC+OC emissions
    Outputs:
        Forcing timeseries
    """

    year, em_SOx, em_BC, em_OC = emissions[:,[0, 5, 9, 10]].T

    def _ERFaci(em,
        ghan_params=np.array([-1.95011431, 0.01107147, 0.01387492])):
        scale = ghan_params[0]
        b_SOx = ghan_params[1]
        b_POM = ghan_params[2]
        return scale*np.log(1+b_SOx*em[0]+b_POM*em[1])

    # PI forcing was not zero as there were some emissions. Use estimates
    # from Skeie et al, 2011 for 1750 forcing.
    E_1765 = np.array([1.0, 11.2])
    nt = len(year)
    F_pd = np.zeros(nt)
    for i in range(nt):
        if year[i]>=1850 or fix_pre1850_RCP==False:
            F_pd[i] = _ERFaci([em_SOx[i], em_BC[i]+em_OC[i]],
                              ghan_params=ghan_params)
        else:
            # linearly interpolate between 1765 and 1850
            E_1850 = np.array([r45e.sox[85], r45e.bc[85]+r45e.oc[85]])
            F_pd[i] = _ERFaci((year[i]-1765)/85.*E_1850 +
                              (1850-year[i])/85.*E_1765,
                              ghan_params=ghan_params)

    # 1765 emissions = zero forcing
    F_1765 = -0.3002836449793625
    F_2011 = -1.5236182344467388

    # are we rescaling to AR5 best estimate with the default parameters?
    if scale_AR5:
        scale=-0.45/(F_2011-F_1765)
    else:
        scale=1.0

    return (F_pd - F_1765) * scale
<|endoftext|>